Modular TTT: 테스트 시 학습을 조합 가능한 모듈로 재해석

Test-Time Training(TTT)은 시퀀스 모델링을 추론 시점에도 내부 학습 규칙으로 fast weights를 업데이트하는 온라인 학습 문제로 바라본다. 이 접근은 현재 문맥에 즉시 적응할 수 있다는 장점이 있지만, 기존의 TTT 변형들은 대개 각각 따로 하드코딩되어 있어 새로운 방법을 설계하거나 각 구성요소의 역할을 구분하기가 쉽지 않았다.

Modular TTT는 이 문제를 모듈형 설계로 풀어낸다. 내부 학습기를 DAG로 표현하고, 다음 요소를 독립적인 설계 차원으로 드러낸다.

• fast-weight network
• loss function
• learning rate
• weight decay
• normalization

또한 train-view forward, train-view backward, causal query-view라는 원시 규칙을 자동으로 조합해, fast-weight 상태 전이까지 포함한 전체 그래프 수준의 TTT 계산을 구성한다. 덕분에 연구자는 개별 모듈을 바꾸면서도 전체 파이프라인을 다시 작성할 필요가 없다.

이 프레임워크의 가장 큰 장점은 체계적인 ablation이 가능하다는 점이다. 논문은 작은 학습률 초기화, weight decay, 단일층 비선형성이 성능 향상에 도움이 된다고 보고한다. 반면 MSE와 내적 손실은 비슷한 수준의 성능을 보였다. 더 깊은 fast-weight network와 정규화는 활성값이 지나치게 커지기 쉬워 성능을 떨어뜨렸고, residual connection과 gating은 측정 가능한 이득이 거의 없었다.

이 결과는 TTT에서 중요한 것이 단순히 더 많은 장치를 쌓는 일이 아니라, 업데이트 규칙과 구조를 얼마나 적절하게 고르느냐에 달려 있음을 보여준다. 다시 말해, 모듈화는 TTT를 더 복잡하게 만드는 도구가 아니라, 불필요한 복잡성을 걷어내는 방법이기도 하다.

저자들은 이러한 통찰을 바탕으로 최적 변형을 410M과 1.45B 파라미터 모델로 학습했고, 100B tokens 규모에서 훈련 손실과 벤치마크 성능이 Gated DeltaNet과 비슷한 수준이라고 밝혔다. 이는 Modular TTT가 개념 정리 차원을 넘어 실제 경쟁력도 갖출 수 있음을 시사한다.

전반적으로 이 연구는 테스트 시 적응을 활용하는 시퀀스 모델을 더 재사용 가능하고 비교 가능하며 디버깅하기 쉬운 형태로 만드는 방향을 제시한다.

출처: Hugging Face Daily Papers