StableGrad: Batch Normalization 없이 역전파 스케일 제어하기

매우 깊은 신경망 (Neural Networks)을 학습시키기 위해서는 깊이에 따른 크기 전파 (Propagation of magnitudes)를 제어해야 합니다. 이러한 제어가 없다면 활성화 값 (Activations)과 그래디언트 (Gradients)가 소실 (Vanish)되거나 폭주 (Explode)할 수 있으며, 최적화 (Optimization)를 실패하게 만드는 불안정한 영역에 진입할 수 있습니다. 현대적인 아키텍처들은 Batch Normalization, 잔차 연결 (Residual connections), 또는 기타 정규화 레이어 (Normalization layers)를 통해 중간 표현 (Intermediate representations)을 반복적으로 재스케일링 (Re-scale)하거나 우회함으로써 이 문제를 완화하곤 합니다. 하지만 이러한 메커니즘이 항상 적절한 것은 아닙니다. 물리 정보 신경망 (Physics-Informed Neural Networks, PINNs)에서 네트워크는 연속적인 물리장 (Physical field)을 나타내며 입력 미분값 (Input derivatives)이 학습 목표를 정의합니다. 이 때문에 배치 의존적 정규화 (Batch-dependent normalization)는 예측된 물리장과 그 미분값에 비국소적 의존성 (Non-local dependencies)을 도입할 수 있어 문제가 됩니다.

우리는 순방향 모델 (Forward model)을 수정하지 않고도 레이어별 가중치-그래디언트 불균형 (Layer-wise weight-gradient imbalances)을 교정하는 옵티마이저 수준의 스케일 제어 메커니즘인 StableGrad를 제안합니다. 정규화가 역전파 (Backpropagation) 이후 및 옵티마이저 업데이트 (Optimizer update) 이전에만 적용되기 때문에, 네트워크 출력, 그 미분값, 그리고 물리적 잔차 (Physical residual)는 변하지 않은 채 유지됩니다. 우리는 이 재스케일링 (Rescaling)에 의해 유도되는 유효 학습 역학 (Effective training dynamics)을 분석하고, 대상 애플리케이션으로서 깊은 PINNs에서 StableGrad를 평가하며, 진단용 스트레스 테스트로서 Batch Normalization이 없는 합성곱 신경망 (Convolutional networks)을 사용합니다.

PINN 벤치마크에서 StableGrad는 매칭된 깊이의 솔루션 정확도 (Solution accuracy)를 향상시키고, 표준 최적화 환경에서 더 깊은 모델을 더 신뢰할 수 있게 만듭니다. Batch Normalization을 제거할 경우 통상적으로 학습 붕괴 (Training collapse)가 발생하는 ResNet 및 EfficientNet 아키텍처에서도, StableGrad는 다른 아키텍처 변경 없이 최적화를 안정화합니다. 이러한 결과는 가중치-그래디언트 스케일에 대한 옵티마이저 수준의 제어가 순방향 정규화 (Forward normalization)를 사용할 수 없거나 원치 않는 상황에서 실용적인 대안이 될 수 있음을 보여줍니다.