평균 모드 스크리밍: 1000개 레이어의 디퓨전 트랜스포머를 위한 평균-분산 분할 잔차 연결

Diffusion Transformer(DiT)를 수백 개 이상의 레이어로 확장할 때 발생하는 갑작스러운 학습 붕괴 현상인 Mean Mode Screaming(MMS)의 원인을 규명했다. 이를 해결하는 MV-Split Residuals 기법을 통해 1000개 레이어의 초거대 모델도 안정적으로 학습할 수 있는 기술적 토대를 마련했다.

Transformer의 핵심인 Self-Attention은 모든 토큰의 관계를 계산하지만, 레이어가 깊어질수록 모든 토큰이 비슷해지는 '표현 균일화' 문제가 발생한다. 특히 Diffusion 모델에서는 각 레이어를 지날 때마다 토큰들의 평균값은 보존되지만, 개별 토큰의 고유한 특징을 담은 분산 성분은 수축되는 기하학적 비대칭성이 존재한다.

이 논문은 학습 도중 특정 시점에 그래디언트가 평균 성분에만 집중적으로 쏠리는 'Mean Mode Screaming' 현상이 발생함을 발견했다. 이때 평균 성분이 네트워크를 장악하면 Softmax 연산의 특성상 쿼리(Query)와 키(Key)를 학습시키는 신호가 완전히 사라져 모델이 더 이상 새로운 특징을 배우지 못하고 붕괴된다.

연구팀은 잔차 연결(Residual Connection)을 통과하는 신호를 '평균'과 '분산'으로 쪼개어 관리하는 아이디어를 냈다. 평균 성분은 너무 커지지 않게 억제(leaky integrator)하고, 실제 이미지의 세부 정보를 담은 분산 성분은 충분히 전달되도록 독립적인 통로를 만들어 줌으로써 1000층이라는 극단적인 깊이에서도 정보가 손실되지 않고 흐를 수 있게 했다.

네트워크의 상태를 시퀀스 평균 성분인 μ(X)와 중심화된 변동 성분인 c(X)로 분해하는 기하학적 프레임워크를 구축했다. Row-stochastic Attention 행렬 A에 대해 Aμ(X) = μ(X)가 성립하여 평균은 보존되지만, 분산 성분은 수축될 수 있음을 보였다.

MV-Split Residuals는 기존의 Post-Norm 구조를 대체한다. 잔차 분기 출력 Fl을 평균 성분 JFl과 분산 성분 PFl로 분리한다. 이후 Zl = Xl + β ⊙ (PFl) + α ⊙ J(Fl - Xl) 연산을 수행한다. 여기서 α와 β는 학습 가능한 벡터이다. [입력값 Xl, Fl → 평균/분산 분리 및 독립 이득 α, β 곱셈 → 결과값 Zl 산출] 과정을 거치며, α를 통해 평균 성분의 누적을 제어하고 β를 통해 특징 학습을 위한 분산 신호를 보존한다.

학습 안정성을 위해 Softmax Jacobian의 특성을 활용했다. 토큰이 균일화되어 Vj = v가 되면 ∂L/∂Si = 0이 되어 Q/K 경로의 학습이 중단됨을 증명했다. 이를 방지하기 위해 잔차 경로에서 평균 성분을 감쇠시키는 leaky trunk-mean replacement 메커니즘을 적용하여 네트워크가 평균 지배 상태에 빠지지 않도록 설계했다.

400개 레이어 DiT 모델 비교 실험에서 기존의 LayerScale 기법보다 빠른 수렴 속도를 보였으며, 일반적인 Post-Norm 모델이 겪는 학습 붕괴를 완벽히 방지했다. ImageNet-2012 데이터셋 기준 FID-50K 점수에서 MV-Split-400L 모델은 2.60을 기록하여 LayerScale-400L의 2.90보다 우수한 성능을 보였다.

1000개 레이어 규모의 DiT 모델(13.64B 파라미터) 실험에서도 학습이 끝까지 안정적으로 유지됨을 확인했다. 해당 모델은 FID 2.77, Inception Score 217.3을 달성하며 초심층 구조에서의 확장성을 입증했다.

Ablation Study를 통해 Attention 분기뿐만 아니라 FFN(Feed-Forward Network) 분기 모두에 MV-Split을 적용해야만 MMS 현상을 완전히 억제할 수 있음을 확인했다. Attention에만 적용했을 경우 그래디언트 스파이크가 FFN 분기로 전이되어 결국 학습이 붕괴되는 현상이 관찰되었다.

MMS는 그래디언트가 평균 성분 방향으로 정렬되면서 O(T) 규모로 증폭되는 현상이다. 논문은 ∇WL = Tδ̄ȳᵀ + Σδ̃tỹtᵀ 식을 통해 그래디언트가 평균 일관성 성분과 중심화된 확산 성분으로 정확히 분해됨을 보였다. [입력 평균 ȳ, 그래디언트 평균 δ̄ → 외적 후 T배 → 평균 성분 그래디언트] 계산을 통해 토큰이 정렬될수록 그래디언트 에너지가 폭발함을 설명한다.

MV-Split은 Trunk의 평균 성분을 (1-α)만큼 수축시킨 후 새로운 잔차의 평균을 더하는 leaky integrator 방식으로 동작한다. 이는 LayerScale이 Trunk의 평균을 그대로 방치하는 것과 대조적이다. 또한 Triton을 이용해 RoPE, QK-Norm, SwiGLU, MV-Split+RMSNorm을 융합(Fusion) 구현하여 메모리 오버헤드를 줄이고 학습 속도를 PyTorch Baseline 대비 2.76배 향상시켰다.

본 연구의 분석은 Transformer 구조, 특히 Softmax Attention의 특성에 의존하고 있어 Mamba나 합성곱 기반의 모델에 직접적으로 적용되는지는 추가 검증이 필요하다. 또한 MMS가 발생하는 정확한 시점(t*)을 사전에 예측하는 수식은 아직 도출하지 못했다.