End-to-End Test-Time Training for Long Context

Overview

• long context에서 attention 계산 비용 문제 해결을 위한 방법 제안
• 기존 접근은 sparse attention, 압축, cache 등 구조 변경 중심
• TTT-E2E는 test-time에 모델을 업데이트하는 새로운 패러다임
• full attention 없이도 long context에서 유사한 성능 달성 목표

Key Takeaways

Problem Setting

• self-attention은 context 길이에 따라 계산량이 급격히 증가
• 기존 방법
  • Longformer, Linformer: attention 구조 변경
  • Transformer-XL: cache 기반 처리
• 한계
  • 구조 변경 시 성능 손실 가능
  • 긴 문맥 전체 정보 활용 어려움
• 핵심 문제: 효율성과 성능을 동시에 만족하는 long context 처리 방법

Main Idea

• test-time에서 모델을 지속적으로 업데이트하여 문맥 정보를 학습
• Test-Time Training
  • next token prediction loss 기반으로 매 step 파라미터 업데이트
  • 문맥 정보를 모델 파라미터에 축적
• TTT-E2E 구조
  • outer loop: 초기 파라미터(meta-learning) 학습
  • inner loop: test-time sequential update
  • 별도 auxiliary loss 없이 NTP loss로 end-to-end 학습
• Mini-batch TTT
  • token 단위 대신 mini-batch 단위 업데이트
  • 안정성 및 효율성 향상
  • sliding window 기반으로 구성
• Sliding Window Attention
  • 제한된 window 내에서 attention 수행
  • window size > batch size로 설정하여 정보 전달 유지
  • long context에서도 안정적 처리 가능
• Efficient Update Strategy
  • MLP layer만 test-time에 업데이트
  • 전체 block 중 일부(약 1/4)만 선택적으로 업데이트
  • 일부 MLP는 고정하여 catastrophic forgetting 완화

Result

• 8K 이상의 context에서 full attention과 유사한 성능 달성
• context 길이가 증가해도 성능 저하 없이 유지
• 32K context까지 확장 가능
• decoding 과정에서도 더 낮은 perplexity 달성
• FLOPs 증가 없이 효율적인 long context 처리