기술 요약: 딥 디퓨전 트랜스포머를 위한 평균 모드 비명 및 MV-분할 잔차

1. 문제 제기

디퓨전 트랜스포머 (DiT) 를 극단적인 깊이 (수백에서 수천 개 레이어) 로 확장하는 것은 표준적인 기울기 폭주 또는 기울기 소실 휴리스틱과 구별되는 구조적 취약점을 도입합니다. 저자들은 토큰 표현이 동질화되고 중심화된 토큰 변동이 억제되는 특정 실패 모드를 **평균 지배적 붕괴 (Mean-Dominated Collapse)**라고 명명하여 식별했습니다. 이는 네트워크가 사소한 평균을 예측하게 만들어, 명백한 NaN 이나 순방향 포화 없이 훈련 손실이 초기화 수준으로 발산하는 상태로 이어집니다.

이 상태로의 급격한 진입 사건은 **평균 모드 비명 (Mean Mode Screaming, MMS)**으로 특징지어집니다. MMS 는 평균-일관성 기울기 성분의 날카로운 급증, 잔차 분기의 급속한 개방, 그리고 그 후의 Query/Key(Q/K) 기울기 억제로 정의됩니다. 이 현상은 수천 스텝 동안 훈련이 안정적으로 보이다가 몇 번의 업데이트 내에 발산할 때에도 발생할 수 있습니다.

기존의 깊이 안정화기 (ReZero 및 LayerScale 등) 는 전체 잔차 분기에 등방성 게이트 (스칼라 또는 채널별) 를 적용합니다. 이러한 방법들은 훈련을 안정화시키지만, 평균 성분과 중심화된 성분을 동시에 억제합니다. 이러한 등방성 감쇠는 공간적으로 변하는 특징 학습에 필요한 중심화된 신호를 감쇠시켜 수렴을 늦춥니다.

2. 방법론: 기작 분석 및 MV-분할

2.1 기작 분석

저자들은 실패 역학을 두 가지 주요 기작으로 분해합니다:

1. 어텐션의 기하학적 비대칭: 행 확률적 어텐션은 순수 평균 상태를 엄격하게 보존하지만 (명제 1), 중심화된 부분공간에서는 수축적으로 작용합니다 (명제 2). 레이어가 깊어질수록 중심화된 성분은 잔차 분기를 통한 보충에 전적으로 의존합니다. 만약 잔차 업데이트가 평균 성분에 의해 지배되면, 네트워크는 순수 평균 상태로 붕괴합니다.
2. 기울기 분해 및 증폭: 저자들은 토큰별 선형 맵 (잔차 작성기) 에 대한 기울기가 **평균-일관성 성분 ($\Delta W_\mu$)**과 **중심화된 성분 ($\Delta W_c$)**으로 정확한 가법적 분해를 허용함을 증명합니다.
   • 평균-일관성 성분은 토큰이 정렬될 때 $O(T)$로 스케일링되며 (여기서 $T$는 시퀀스 길이), 일관성 누적 영역에 진입합니다.
   • 중심화된 성분은 확산적으로 스케일링됩니다.
   • MMS 유발: 토큰 표현과 어조이트가 동질화됨에 따라, 기울기 확장 내의 부호 있는 비대각 항들이 상쇄되지 않습니다. 이로 인해 평균-일관성 기울기가 급증하여 잔차 업데이트를 지배합니다.
3. Q/K 기울기 소멸: 값들이 동질화되면, Softmax 야코비안의 영공간이 어텐션 가중 기울기의 상수 성분을 억제합니다. 이는 Q/K 기울기가 소멸하게 하여 (보조 1), 어텐션 메커니즘이 더 이상 토큰 변동을 복원할 수 없기 때문에 네트워크를 붕괴된 상태에 고정시킵니다.

2.2 제안된 해결책: 평균 - 분산 분할 (MV-Split) 잔차

이를 해결하기 위해 저자들은 평균과 중심화된 부분공간의 제어를 분리하는 **MV-분할 잔차 (MV-Split Residuals)**를 제안합니다. 단일 등방성 이득 대신, MV-분할은 부분공간 라우팅 병합을 사용합니다:

$$Z_l = X_l + \beta \odot (P F_l) + \alpha \odot J(F_l - X_l)$$

여기서:

• $X_l$은 트렁크 입력이고, $F_l$은 분기 출력입니다.
• $J$는 시퀀스 평균 부분공간으로의 프로젝터이며, $P = I - J$는 중심화된 부분공간으로의 프로젝터입니다.
• $\alpha$와 $\beta$는 토큰 전체에 브로드캐스트되는 학습 가능한 벡터입니다.

주요 역학:

• 중심화된 경로: 중심화된 성분은 공간적으로 변하는 특징을 학습할 수 있는 능력을 보존하는 이득 $\beta$를 가진 표준 잔차 업데이트를 받습니다.
• 평균 경로: 평균 성분은 "누출성 트렁크 평균 교체"를 겪습니다. 트렁크 평균은 새로운 보정이 추가되기 전에 $(1-\alpha)$로 수축됩니다. 이는 중심화된 경로를 축소하지 않으면서 평균-일관성 업데이트의 축적을 방지합니다.
• 역방향 역학: 기울기 흐름도 유사하게 분할되어, 평균-일관성 기울기 성분 ($\Delta W_\mu$) 이 $\alpha$에 의해 감쇠되도록 하고, 이는 $\beta$에 의해 제어되는 중심화된 기울기 성분 ($\Delta W_c$) 에 영향을 주지 않습니다.

3. 주요 기여

1. 붕괴의 특성화: 이 논문은 평균 지배적 붕괴 상태와 급격한 진입 사건인 MMS 를 구분합니다. 표준 초기화는 점진적인 깊이별 붕괴로 이어지는 반면, 0 으로 초기화된 작성기는 지연된 MMS 급증을 유발함을 보여줍니다.
2. 기작적 통찰: 저자들은 행 확률적 어텐션이 순수 평균 상태를 보존하면서 중심화된 부분공간을 수축한다는 이론적 증명을 제공합니다. 그들은 토큰 표현과 어조이트의 정렬에 따라 일관성 기울기 누적로의 전환을 연결하는 **정렬 - 증폭 법칙 (Alignment-Amplification Law)**을 유도합니다. 또한, 값의 동질화가 Softmax 야코비안 영공간을 통해 Q/K 학습을 구조적으로 억제함을 증명합니다.
3. 방법 및 검증: 평균 경로를 감쇠시키면서 중심화된 경로를 보존하여 훈련을 안정화시키는 MV-분할 잔차 제안.
   • 400 레이어 검증: 매칭된 실험에서 MV-분할은 안정화되지 않은 베이스라인에서 관찰된 발산적 붕괴를 방지하고, LayerScale 보다 훨씬 빠르게 수렴합니다.
   • 1000 레이어 규모 검증: 저자들은 MV-분할을 사용하여 1000 레이어 DiT 를 성공적으로 훈련시켰으며, 표준 방법이 실패하는 극단적인 깊이에서도 아키텍처가 안정적으로 훈련 가능함을 입증했습니다.

4. 실험 결과

• 안정성: ImageNet-2012 잠재 공간에서의 400 레이어 실행에서, 안정화되지 않은 베이스라인과 감소된 학습률 베이스라인 모두 발산 (붕괴) 했습니다. LayerScale 은 안정적이었지만 수렴 속도가 느렸습니다. MV-분할은 안정적이었으며 비발산 실행 중 최고의 성능을 달성했습니다.
• 성능: 5 만 훈련 스텝에서 400 레이어 MV-분할 모델은 FID-50K 2.60 및 Inception Score 185.5를 달성하여 LayerScale(FID 2.90, IS 165.5) 을 능가했습니다.
• 기울기 역학: 작성기 기울기 분석은 MV-분할이 평균-일관성 성분 ($G_{mean}$) 을 경계 짓는 동시에 중심화된 성분 ($G_{ctr}$) 에 대해 더 높고 안정적인 밴드를 유지함을 보여주었습니다. 반면, LayerScale 은 두 성분을 동일하게 억제했습니다.
• 1000 레이어 실행: 동일한 MV-분할 설계를 사용한 별도의 1000 레이어 실행은 훈련 기간 내내 안정적으로 유지되었으며, 5 만 스텝에서 FID 2.77 및 IS 217.3 에 도달했습니다.
• 정성적 결과: 1000 레이어 모델은 자연어 캡션에 조건부인 다양한 의미 범주 (예: 동물, 사물, 풍경) 에 걸쳐 고충실도 텍스트 - 이미지 샘플을 생성했습니다.

5. 중요성 및 주장

이 논문은 MV-분할이 초심층 Post-Norm DiT 에서 이전까지 다루어지지 않았던 특정 실패 경로, 즉 $O(T)$로 성장하는 평균-일관성 작성기 업데이트와 수축적인 중심화된 경로 사이의 불균형을 해결한다고 주장합니다.

• 아키텍처적 필요성: 저자들은 등방성 게이트 (LayerScale 과 같은) 가 특징 학습에 필요한 신호 (중심화된 변동) 를 감쇠시키기 때문에 극단적인 깊이에는 불충분하다고 주장합니다. MV-분할은 필요한 이방성 제어를 제공합니다.
• 확장성: 1000 레이어 DiT 의 성공적인 훈련은 경계 규모 검증으로 기능하여, 잔차 인터페이스가 적절히 관리된다면 아키텍처를 극단적인 깊이로 확장할 수 있음을 확립합니다.
• 기작적 이해: 이 연구는 휴리스틱 관찰을 넘어 기울기 모드 및 어텐션 역학의 분해로 나아가는, 심층 DiT 가 붕괴하는 이유에 대한 엄격한 기하학적 및 대수적 설명을 제공합니다.

저자들은 그들의 분석이 Softmax 어텐션 및 Post-Norm 아키텍처에 구체적임을 지적합니다. 그들은 MMS 가 모든 심층 트랜스포머 불안정성을 설명한다고 주장하지는 않지만, 오히려 MV-분할이 잔차 인터페이스에서 해결하는 특정 잔차 부분공간 실패 경로를 식별한다고 명시합니다.