이중 사전 조건부 최적화 (DoPr): 검증 손실이 아닌 테스트 시간 성능을 위한 최적화

Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss

최근의 딥러닝 응용 분야에서 신경망은 종종 단일 단계 예측 손실(예: $L^2$ 회귀, 교차 엔트로피)을 사용하여 학습되지만, 학습된 모델은 자체 예측에 따라 동작합니다. 이러한 예로는 자기 회귀 언어 모델링, 흐름 기반 생성 모델링, 로봇 정책 학습 등이 있습니다. 이러한 설정에서는 '테스트 시간 피드백(TTF)'이라는 현상이 발생하는데, 이는 학습/검증 손실과 실제 성능 지표(예: 작업 성공률, 생성 품질) 간의 불일치를 초래하며, 작업 길이가 증가함에 따라 이 불일치는 더욱 심화됩니다. TTF 환경에서 이러한 훈련-테스트 격차 문제를 해결하기 위해 데이터 정제, 아키텍처 설계 및 목적 함수 설계를 제안해 왔지만, 본 논문에서는 오류 누적을 완화하기 위한 새로운 설계 축으로 최적화를 제시합니다. 특히, TTF의 어려움에 특화된 새로운 최적화 패러다임인 '이중 사전 조건부 최적화(DoPr)'를 소개합니다. DoPr은 Adam 및 Muon과 같은 기울기 기반 사전 조건부 기법과 KFAC와 같은 활성화 함수 기반 사전 조건부(AP) 기법을 결합합니다. 실험 결과, AP를 추가하면 다양한 TTF 환경에서 모델의 성능 향상을 가져오는 효과적인 방법임을 확인했습니다. 흥미롭게도, 이러한 테스트 시간 성능 개선은 검증 손실의 일관된 개선과 반드시 동반되지 않으며, 이는 단일 단계 감독 학습 목표로 훈련된 모델을 적절하게 평가하는 방법에 대한 새로운 질문을 제기합니다.

Many modern applications of deep learning involve training a neural network via a one-step prediction loss (e.g., $L^2$ regression, cross-entropy), but deploy the network by rolling out along its own predictions. Key examples include autoregressive language modeling, flow-based generative modeling, and robot policy learning. It is well-documented that these settings induce a phenomenon we call test-time feedback (TTF): the mismatch between the training/validation loss and downstream metrics of interest, such as task success rate and generation quality, which grows with task length. While data curation, architecture, and objective design have been proposed to combat train-test shift in TTF settings, this paper proposes optimization as a new design axis to mitigate error accumulation. Specifically, we introduce a new optimization paradigm called double-preconditioning (DoPr) uniquely tailored to the challenges of TTF. DoPr combines gradient-wise preconditioning, as in Adam and Muon, with activation-wise preconditioning (AP), such as in KFAC. We show that the addition of AP yields a drop-in intervention for increasing downstream model performance across a range of TTF settings. Interestingly, these gains in test-time performance do not consistently accompany improvements in validation loss, opening new questions about how to properly evaluate models trained with one-step supervised objectives.