FlexAct: Why Learn when you can Pick?

기술 요약: Flex-Act

문제 정의
활성화 함수는 심층 신경망 내의 신호 전파, 학습 가능성 및 일반화에 결정적인 역할을 수행함에도 불구하고, 일반적으로 모든 레이어와 태스크에 균일하게 적용되는 고정된 하이퍼파라미터(예: ReLU, GELU, Swish)로 취급됩니다. 이러한 정적 패러다임은 단일 비선형 함수가 레이어 깊이나 데이터 레짐에 관계없이 최적이라는 가정을 전제로 하며, 이는 잠재적으로 최적의 유도 편향(inductive bias) 상실, 그래디언트 파쇄(shattered gradients), 또는 포화(saturation) 문제를 야기할 수 있습니다. 매개변수화된 활성화 함수(예: PReLU, SPLASH)나 학습 가능한 함수 근사치들이 존재하기는 하지만, 이들은 종종 높은 최적화 복잡성을 도입하거나, 해석력을 감소시키거나, 특정 함수군(functional family)을 반드시 선택해야 하는 제약을 가집니다. 저자들은 단일 비선형 함수가 보편적으로 최적일 수는 없으며, 심층 신경망이 수동적인 하이퍼파라미터 튜닝 없이 각 레이어에 가장 적합한 활성화 함수를 동적으로 선택할 수 있는 능력을 갖추었을 때 더 큰 이득을 얻을 것이라고 주장합니다.

방법론
본 논문은 훈련 과정 중 사전 정의된 집합으로부터 레이어별로 이산적인 활성화 함수를 선택할 수 있게 하는 프레임워크인 Flex-Act를 소개합니다. 핵심 메커니즘은 이산적 선택을 미분 가능한 완화(relaxation)로 만드는 Gumbel-Softmax 트릭에 기반합니다.

1. Gumbel-Softmax를 통한 이산 라우팅:
   각 레이어 $i$에 대해, 모델은 아핀 변환(affine transformation) $h_i$에 적용되는 후보 활성화 함수들의 볼록 결합(convex combination) $\{\sigma^{(1)}, \dots, \sigma^{(p)}\}$을 계산합니다. 선택 확률 $p_i$는 Gumbel-Softmax 분포를 사용하여 학습 가능한 로짓(logits) $\pi_i$로부터 도출됩니다:
   $$X_i = \sum_{j=1}^{p} p_i^{(j)} \cdot \sigma^{(j)}(h_i)$$
   여기서 $p_i \sim \text{GumbelSoftmax}(\log \pi_i, \tau)$입니다. 온도 $\tau$가 점진적으로 감소(annealing)함에 따라, 모델은 혼합된 활성화 함수를 탐색하는 단계에서 특정 이산 함수로 확정하는 단계로 전이합니다. 이를 통해 선택 과정을 엔드 투 엔드(end-to-end) 방식으로 경사 하강법을 통해 최적화할 수 있습니다.

2. 그래디언트 정규화를 통한 편향 수정:
   저자들은 단순한 Gumbel-Softmax 라우팅이 가진 치명적인 결함, 즉 무제한 함수(ReLU와 같은)로 향하는 스케일 유도 편향(scale-induced bias)을 식별했습니다. 역전파는 활성화 크기에 따라 그래디언트 크기를 조절하기 때문에, 데이터의 적합성과 관계없이 출력값이 큰 함수가 선택 과정을 지배하게 됩니다.
   이를 해결하기 위해 Flex-Act는 **그래디언트 노름 기반 정규화 도구(gradient-norm-based regularizer)**를 도입합니다. 이 도구는 각 후보 활성화 함수의 입력에 대한 그래디언트 노름을 계산하고, 이를 의사 확률(pseudo-probabilities)로 변환합니다. 전체 손실 함수는 기본 태스크 손실($L_{task}$)과 학습된 라우팅 확률을 이러한 그래디언트 기반 의사 레이블과 일치시키는 KL 발산 항($L_{KL}$)의 결합으로 구성됩니다:
   $$L = L_{task} + \alpha L_{KL}$$
   이 정규화는 모델이 단순히 크기 차이 때문에 무제한 함수로 편향되는 것을 방지하며, 선택이 기능적 적합성에 의해 결정되도록 보장합니다.

주요 기여

• 이산 활성화 라우팅: 엔드 투 엔드 최적화를 가능하게 하기 위해 Gumbel-Softmax 재매개변수화(reparameterization)를 사용하는 새로운 메커니즘을 제시합니다.
• 분석적 그래디언트 유도: 이산 활성화 라우팅을 사용하는 네트워크를 위한 역전파 방정식을 공식적으로 유도하여, 선택 과정의 비미분성 문제를 해결했습니다.
• 그래디언트 노름 정규화: 이산 라우팅에서 발생하는 무제한 활성화 함수로의 내재적 편향을 교정하기 위한 특정 기술을 제공하며, 이는 선택 정확도를 크게 향상시킵니다.
• 해석력 및 모듈성: 활성화 선택을 이산적이고 학습 가능한 변수로 취급함으로써, 서로 다른 깊이에서 어떤 비선형성이 사용되는지에 대한 통찰을 제공하며, 매개변수화된 함수군에 대한 모듈형 대안을 제시합니다.

실험 결과
논문은 합성 회귀 태스크와 실제 벤치마크 모두에서 Flex-Act를 평가합니다.

• 합성 회귀 (Synthetic Regression): 지면 진리(ground-truth) 활성화 함수(ReLU, Sigmoid, Tanh, LeakyReLU, Identity)가 알려진 태스크에서, Flex-Act는 일관되게 올바른 함수로 수렴합니다.
  • 정규화가 없는 경우 ($\alpha=0$), 모델은 그래디언트 크기 편향으로 인해 렐루(ReLU)로 귀결되어 유계 함수(Sigmoid/Tanh)를 선택하는 데 실패합니다.
  • 정규화가 적용된 경우 ($\alpha=0.3$), Flex-Act는 모든 지면 진리 시나리오에서 거의 0에 가까운 평균 제곱 오차(MSE)를 달성하며, 여러 번의 훈련이나 활성화 탐색 없이도 고정된 활성화 베이스라인과 대등하거나 이를 능가하는 성능을 보입니다.
• 실제 벤치마크 (Real-World Benchmarks):
  • 이미지 분류 (CIFAR-10/100): ResNet-18 및 ResNet-34에 통합된 Flex-Act(모든 레이어 또는 마지막 레이어 적용)는 고정된 ReLU 베이스라인 및 기타 매개변수화된 방법(PReLU, Swish, AReLU)보다 통계적으로 유의미한 개선을 달성했습니다. 예를 들어, 모든 레이어에 Flex-Act를 적용한 ResNet-34는 CIFAR-10에서 95.75%, CIFAR-100에서 79.42%의 정확도에 도달하여 ReLU 베이스라인을 능가했습니다.
  • 텍스트 분류: BERT-base 모델을 사용한 GLUE 벤치마크에서도 테스트되었으며, 트랜스포머 아키텍처와의 호환성을 입증했습니다.

의의 및 주장
본 논문은 Flex-Act가 적응형 딥러닝을 위한 "플러그 앤 플레이(plug-and-play)" 솔루션을 제공한다고 주장합니다. 저자들은 네트워크가 특정 태스크나 레이어에 적합한 비선형성을 동적으로 복구할 수 있게 함으로써, 실무자들이 수동적인 하이퍼파라미터 탐색이나 다양한 고정 활성화 함수를 사용한 여러 모델의 훈련에 드는 계산 비용을 피할 수 있다고 설명합니다. 이 프레임워크는 이론적 유연성과 실용적 효용 사이의 간극을 메우며, 얕은 합성 태스크부터 심층적인 실제 아키텍처에 이르기까지 확장 가능한 견고하고 해석 가능하며 모듈화된 비선형 설계 방식을 제공합니다. 저자들은 이 방법이 기존의 잘 정의된 함수 형태를 지능적으로 선택함으로써, 파라미터 수를 크게 늘리지 않고도 이러한 이득을 달성한다는 점을 강조합니다.