기술 요약: 이중 전처리 (Double Preconditioning, DoPr)

1. 문제 정의: 테스트 타임 피드백 (TTF) 및 손실-성능 간극

본 논문은 모델이 일단계 지도 학습 목적 함수(예: 다음 토큰 예측, 단계별 모방 학습)로 훈련되지만, 실제 배포 시에는 자신의 예측을 따라 순차적으로 롤아웃(rollout)되는 현대 딥러닝 애플리케이션의 결정적인 괴리를 다룹니다. 저자들은 이 현상을 **테스트 타임 피드백 (Test-Time Feedback, TTF)**이라고 명명합니다.

자기회귀 언어 모델링, 흐름 기반 생성 모델링, 로봇 정책 학습과 같은 TTF 설정에서는 모델의 예측에 포함된 작은 오류가 시간이 지남에 따라 누적됩니다. 이는 배포 중에 마주치는 상태(state)의 분포가 훈련 분포로부터 멀어지게 만드는 현상(TTF shift)을 유발합니다.

식별된 핵심 문제는 표준 옵티마이저(예: Adam, Muon, Shampoo)가 훈련 분포 하에서의 검증 손실($L_{val}$)을 최소화하도록 설계되었다는 점입니다. 그러나 저자들은 $L_{val}$을 최소화하는 것이 반드시 다운스트림 태스크 보상($R_{test}$)이나 모델 자신의 배포 분포 상에서 정의된 이상적인 손실($L_{ideal}$)을 최소화하는 것을 의미하지는 않는다고 주장합니다. 구체적으로, 표준 옵티마이저는 모델이 훈련 분포에는 최적이지만 TTF의 피드백 역학 하에서 일반화에 실패하는 특징(feature)을 학습하게 되는 **불완전한 특징 학습(imperfect feature learning)**을 유도할 수 있으며, 이는 검증 손실이 잘 수렴하더라도 다운스트림 성능 저하로 이어집니다.

2. 방법론: 이중 전처리 (Double Preconditioning, DoPr)

데이터, 아키텍처, 또는 훈련 목적 함수를 수정하지 않고 TTF shift를 완화하기 위해, 저자들은 **이중 전처리 (DoPr)**를 제안합니다. DoPr는 두 가지 서로 다른 전처리 전략을 결ни합한 "플러그인(plug-in)" 최적화 프레임워크입니다.

1. 활성화 전처리 (Activation Preconditioning, AP): 레이어별 활성화의 경험적 공분산($\hat{\Sigma}_z \approx E[zz^\top]$)을 사용하여 그래디언트를 전처리하는 레이어 단위 전처리 기법입니다.

   • 메커니즘: 원시 그래디언트 $G$에 활성화 공분산 행렬의 역행렬을 곱합니다: $M = G \cdot \hat{\Sigma}_z^{-1}$.
   • 이론적 동기: 저자들은 레이어 활성화가 비등방성(non-isotropic, ill-conditioned)일 때 표준 경사 하강법이 편향된 특징 학습을 겪는다는 것을 입증했습니다. AP는 이러한 편향을 교정하여, 상태 공간의 모든 방향에 대해 **균등한 특징 학습(uniform feature learning)**을 촉진합니다. 이는 모델이 훈련 손실에 의해 가중치가 낮게 책정되었더라도 TTF 하에서 오류 누적을 방지하는 데 중요한 방향들에 대해서도 강건한 특징을 학습하도록 보장합니다.
   • 불변성 원리: 본 논문은 AP가 입력 분포의 아핀 변환(affine transformation)에 대한 불변성을 유도함을 입증합니다. 입력 분포가 아핀 변환되더라도, AP 업데이트 규칙은 레이어 출력이 불변하도록 하여 분포 변화에 대한 최적화 궤적을 안정화합니다.

2. 그래디언트 전처리 (Gradient Preconditioning, GP): AP로 전처리된 그래디언트에 적용되는 표준 그래디언트 기반 전처리 기법(예: Adam, Muon, Signum)입니다.

   • 역할: AP가 특징 학습의 기하학적 구조를 다룬다면, GP는 이를 통해 발생하는 수치적 불안정성을 해결합니다. GP 구성 요소는 훈련 과정을 안정화하고, 수렴을 가속화하며, 모멘텀을 제공합니다.
   • 통합: DoPr는 먼저 그래디언트에 AP를 적용한 후, 그 결과를 GP에 전달합니다. 이를 통해 DoPr는 확립된 GP(Adam 또는 Muon 등)의 안정성과 하이퍼파라미터 스케일링 규칙을 활용하면서도, AP의 TTF 완화 특성을 상속받을 수 있습니다.

알고리즘 흐름:
가중치 $W$와 활성화 $z$를 가진 레이어에 대하여:

1. 배치 그래디언트 $G = \nabla_W \mathcal{L}$를 계산한다.
2. 경험적 활성화 공분산 $\hat{\Sigma}_z$를 계산한다.
3. AP 적용: $M = G \cdot (\hat{\Sigma}_z + \gamma \text{tr}(\hat{\Sigma}_z)I)^{-1}$.
4. GP 적용: $D = \text{GP}(M, \dots)$.
5. 가중치 업데이트: $W \leftarrow W - \eta D$.

3. 주요 기여

• 이론적 통찰: 본 논문은 TTF shift와 불완전한 특징 학습 사이의 관계를 엄밀하게 연결합니다. 저자들은 표준 옵티마이저가 낮은 검증 손실을 갖더라도 특징의 부공간 정렬(subspace alignment)이 불량한 솔루션으로 수렴할 수 있으며, 이것이 TTF를 악화시킨다는 것을 증명했습니다 (Proposition 3.2 및 4.1을 통해).
• DoPr 프레임워크: 방향성(AP를 통해 균등한 특징 학습을 보장)과 적응성/안정성(GP를 통해 처리)을 분리하는 새로운 최적화 패러다임을 도입했습니다.
• 드롭인 호환성 (Drop-in Compatibility): DoPr는 기존 옵티마이저의 드롭인 대체제로 설계되었습니다. 모델 아키텍처나 데이터 큐레이션을 변경할 필요가 없습니다.
• 하이퍼파라미터 전이: 저자들은 DoPr가 GP 구성 요소를 유지하기 때문에, Adam이나 Muon과 같은 기본 GP에서 파생된 기존의 하이퍼파라미터 스케일링 규칙(예: $\mu$P 스케일링)이 DoPr로 직접 전이됨을 보여줍니다. 이는 대규모 훈련을 용이하게 합니다.

4. 실험 결과

저자들은 연속 제어, 로보틱스, 언어 모델링 태스크에 걸쳐 DoPr를 평가했습니다. 주요 결과는 다음과 같습니다.

• 연속 제어 (Humanoid-v5): DoPr는 베이스라인 옵티마이저(AdamW, Muon, Signum, AdaMuon) 대비 터미널 보상(terminal reward)을 일관되게 향상시킵니다. 주목할 점은, 이러한 다운스트림 성능의 향상이 훈련 또는 검증 손실의 개선과 반드시 일치하지 않는다는 것입니다. 일부 경우, DoPr는 동일하거나 약간 더 나쁜 손실 수렴도를 보이면서도 더 높은 보상을 달성했습니다.
• 로봇 정책 학습 (Robomimic): "Tool-Hang" 및 "Transport"와 같은 도전적인 태스크에서 DoPr 변형 모델들은 베이스라인보다 높은 성공률을 보였습니다. 이 역시 태스크 성공의 이득이 훈련 손실의 균일한 개선을 동반하지는 않았습니다.
• 언어 모델 (수학적 추론):
  • 3B SFT: DoPr-AdamW는 토큰 수준의 훈련 손실이 비슷하거나 더 높음에도 불구하고, 다양한 학습률에 대해 AdamW보다 GSM8K에서 더 높은 피크 정확도를 달s했습니다.
  • 8B SFT: 대규모 미세 조정(fine-tuning)에서 DoPr-AdamW는 베이스라인 AdamW에 비해 검증 손실과 다운스트림 성능(GSM8K, MATH-500) 사이의 상관관계를 더 강력하게 보여주었습니다. 베이스라인 AdamW 모델은 높은 학습률에서 검증 손실은 개선되었으나 태스크 성능은 저하되는 양상을 보인 반면, DoPr는 견고한 성능을 유지했습니다.
  • 교차 태스크 갈등 (Cross-Task Conflict): DoPr는 태스크 간의 갈등을 줄입니다. GSM8K-CoT에 대한 최적 학습률이 MATH-500에서도 최적에 가까운 값을 가졌는데, 이는 베이스라인 AdamW에서는 관찰되지 않은 특성입니다.

5. 의의 및 주장

본 논문은 최적화가 TTF 환경에서 오류 누적을 완화하기 위한 새로운 설계 축이며, 이는 데이터 큐레이션이나 아키텍처 설계와는 별개의 영역이라고 주장합니다.

• 손실과 성능의 분리: 핵심적인 의의는 TTF 환경에서 다운스트림 성능을 개선하는 것이 반드시 훈련 손실의 개선을 요구하지 않으며(심지어 동반하지 않을 수도 있음) 오히려 그 반대일 수 있음을 입증했다는 점입니다. 이는 훈련 손실 최소화가 모델 품질의 유일한 대리 지표라는 전통적인 관념에 도전합니다.
• 레버로서의 특징 학습: 이 연구는 (AP에 의해 보장되는) "균등한 특징 학습"이 손실 수렴 속도보다 TTF 하에서의 강건성을 위한 결정적인 요소라고 상정합니다.
• 실용적 유용성: DoPr는 모델이 자신의 예측을 따라 롤아웃해야 하는 실제 배포 시나리오에서 모델의 강건성을 향상시키는 실용적이고 오버헤드가 낮은 개입(plug-in) 수단을 제공합니다.

저자들은 "옵티마이저의 프리러치(free-lunch)"가 여전히 존재하는지 질문하며, 옵티마이저의 방향성과 적응성을 분리하는 것이 특히 훈련 목적 함수가 배포 현실과 괴리되는 태스크에서 중요한 연구 분야임을 시사하며 마무리합니다.