AI바라기의 인공지능

LLM : 빠른 논문 리뷰 : End-to-End Test-Time Training for Long Context 본문

논문리뷰

LLM : 빠른 논문 리뷰 : End-to-End Test-Time Training for Long Context

AI바라기 2026. 1. 16. 16:43

용어 설명 (Vocabulary)

  • Test-Time Training (TTT): 모델을 학습시킨 후 고정해서 사용하는 것이 아니라, **Inference(추론) 단계에서도 입력 데이터(Context)를 이용해 모델을 실시간으로 학습(업데이트)**시키는 기법입니다.
  • Sliding-Window Attention (SWA): 전체 문맥을 다 보는 것이 아니라, **최근의 특정 길이(Window size)**만큼의 토큰들에 대해서만 Attention을 수행하는 방식입니다. 계산 효율성은 높지만 먼 과거의 정보를 잊어버린다는 단점이 있습니다.
  • Full Attention: Transformer의 기본 방식으로, 현재 토큰이 과거의 모든 토큰과 상호작용합니다. 성능은 좋지만 문맥 길이가 길어질수록 계산 비용이 제곱으로 증가합니다.
  • Meta-learning: '학습하는 법을 학습하는 것'입니다. 이 논문에서는 TTT가 추론 시점에 빠르게 적응할 수 있도록, 초기 가중치($W_0$) 자체를 최적화하는 과정을 의미합니다.
  • Inner Loop / Outer Loop:
    • Inner Loop: 추론 시점에 입력된 Context를 보고 모델(MLP)을 업데이트하는 과정.
    • Outer Loop: Inner Loop가 잘 작동하도록 모델의 초기 상태(Initialization)를 학습시키는 전체 훈련 과정.
  • Needle in a Haystack (NIAH): 긴 텍스트(건초더미) 속에 숨겨진 특정 정보(바늘)를 정확히 찾아내는 능력을 평가하는 테스트입니다.

Purpose of the Paper

  • 기존 연구의 한계 (Efficiency vs. Performance Trade-off):
    • Transformer with Full Attention: 긴 문맥(Long Context)에서 성능은 가장 좋지만, 계산 비용이 문맥 길이에 따라 제곱(quadratic)으로 증가하여 확장이 어렵습니다.
    • RNNs (Mamba, Gated DeltaNet 등): 계산 비용은 일정(constant)하지만, 문맥이 길어질수록 Full Attention만큼 성능이 나오지 않습니다(정보 소실).
  • 새로운 접근 방식 (Continual Learning):
    • 이 논문은 긴 문맥 처리를 아키텍처 설계의 문제가 아닌 Continual Learning(지속적 학습)의 문제로 재정의합니다.
    • 인간이 모든 과거를 완벽하게 기억(Recall)하지 않고도 경험을 통해 직관을 **압축(Compress)**하여 학습하듯이, 언어 모델도 문맥을 KV Cache에 저장하는 대신 모델의 파라미터(Weights)에 압축하여 업데이트해야 한다고 주장합니다.

Key Contributions & Novelty

  • TTT-E2E (End-to-End Test-Time Training) 제안:
    • Standard Architecture: 복잡한 커스텀 레이어 대신, 표준 Transformer 구조에 **Sliding-Window Attention (SWA)**을 적용하고, MLP(Multi-Layer Perceptron) 레이어를 TTT 대상으로 삼았습니다.
    • Integration: 추론 시점(Test time)에 입력된 긴 문맥을 바탕으로 Next-token prediction을 수행하여 MLP 레이어의 가중치를 즉시 업데이트합니다.
    • Meta-learning Initialization: 단순히 TTT를 적용하는 것을 넘어, 훈련 시점(Training time)에 Meta-learning을 통해 TTT가 가장 잘 작동할 수 있는 최적의 초기 가중치를 학습합니다.
  • 기존 TTT 연구 대비 차별점 (Novelty):
    • 기존의 TTT 연구들(예: TTT-KVB)은 Self-Attention의 메커니즘을 모방하기 위해 복잡한 손실 함수나 구조를 사용했습니다.
    • 반면, TTT-E2E는 모델의 최종 목표인 Next-token prediction loss를 Inner Loop(추론 업데이트)와 Outer Loop(전체 학습)에 동일하게 사용하는 End-to-End 방식을 채택하여 구조를 단순화하면서도 성능을 극대화했습니다.

Experimental Highlights

  • State-of-the-art Scalability:
    • 3B 파라미터 모델과 164B 토큰으로 훈련된 실험에서, TTT-E2E는 문맥 길이가 길어질수록 Transformer with Full Attention과 동일한 수준의 Loss 감소(성능 향상) 추세를 보였습니다.
    • 반면, 주요 경쟁 모델인 Mamba 2나 Gated DeltaNet은 긴 문맥(Long context)에서 성능 향상 폭이 둔화되거나 멈췄습니다.
  • Inference Speed & Latency:
    • 128k 길이의 문맥에서 TTT-E2E는 Full Attention 대비 2.7배 더 빠른 추론 속도를 기록했습니다.
    • RNN 계열처럼 문맥 길이에 상관없이 **Constant Inference Latency(상수 시간 지연)**를 달성했습니다.
  • Unique Feature:
    • TTT-E2E는 문맥의 초반부 토큰들(early tokens)에 대해서도 Full Attention보다 더 낮은 Perplexity(더 좋은 예측)를 보였습니다. 이는 TTT가 "미래의 모든 가능성"을 대비하는 것이 아니라, "현재의 문맥"에 집중하도록 가중치를 최적화하기 때문입니다.

Limitations and Future Work

  • Training Latency (치명적 한계):
    • Limitations: 학습 시 'Gradient의 Gradient(이중 미분)'를 계산해야 하므로, 8k 문맥 기준으로 Full Attention보다 학습 속도가 약 3.4배 느립니다. 또한, 현재 구현체는 FlashAttention과 같은 하드웨어 가속 커널을 훈련 단계에서 활용하지 못합니다.
    • Future Work: 커스텀 Attention 커널 개발 및 사전 학습된(Pre-trained) Transformer에서 TTT로 전이 학습(Initialize)하는 방식을 통해 학습 비용을 절감할 계획입니다.
  • Exact Recall (Needle in a Haystack):
    • Limitations: TTT-E2E는 정보를 가중치에 '압축'하는 방식이므로, Full Attention처럼 과거의 정보를 손실 없이 완벽하게 불러오는(Recall) 능력은 떨어집니다. Needle in a Haystack 테스트에서 Full Attention보다 낮은 성능을 보였습니다.
    • 의의: 이 논문은 "모든 것을 기억하는 것"보다 "문맥을 이해하고 일반화하는 것"이 언어 모델링에 더 중요함을 시사합니다.

Overall Summary

이 논문은 긴 문맥을 처리하기 위해 과거 정보를 캐싱하는 대신, 추론 시점에 모델이 스스로 학습(TTT)하여 정보를 가중치에 압축하는 TTT-E2E 방법론을 제안합니다. 실험 결과, TTT-E2E는 Full Attention 모델의 우수한 성능 확장성을 유지하면서도 RNN 모델의 빠른 추론 속도라는 장점을 모두 확보했습니다. 이는 거대 언어 모델(LLM)이 수백만 토큰 이상의 초장문 문맥을 처리할 때 겪는 '메모리 및 속도 병목' 문제를 해결할 수 있는 강력한 대안으로, 아키텍처 중심의 연구에서 학습 알고리즘 중심의 연구로의 패러다임 전환을 제시합니다.

쉬운 설명 (Easy Explanation)

Full Attention 모델이 "시험 칠 때 오픈북으로 교과서의 모든 페이지를 일일이 다시 찾아보는 학생"이라면,

TTT-E2E 모델은 "수업을 들으면서 내용을 완전히 이해해서 머릿속(가중치)에 넣고, 시험 칠 때는 책을 안 보고 바로 답을 적는 똑똑한 학생"과 같습니다.

  • 오픈북 학생(Full Attention)은 토시 하나 안 틀리고 기억해낼 수 있지만(Recall 우수), 책이 두꺼워지면 답을 찾는 데 너무 오래 걸립니다(느린 속도, 비용 증가).
  • 이해하고 푼 학생(TTT-E2E)은 교수가 했던 농담까지 완벽히 기억하진 못할 수 있어도(Recall 약점), 핵심 내용은 완벽히 파악하고 있으며 문제의 양이 많아져도 막힘없이 빠르게 풀어냅니다(빠른 속도, 성능 유지).

이 논문은 AI가 무조건 모든 데이터를 짊어지고 다니는 대신, **실시간으로 공부(Test-Time Training)**하게 만들면 훨씬 효율적이라는 것을 증명한 연구입니다.

 

 

 

Abstract

우리는 long-context language modeling을 architecture 디자인보다는 continual learning의 문제로 공식화합니다. 이러한 공식화 하에서, 우리는 standard architecture인 sliding-window attention을 갖춘 Transformer만을 사용합니다.

그러나 우리의 model은 주어진 context에 대한 next-token prediction을 통해 test time에 지속적으로 학습하며, 읽어들인 context를 weights로 압축합니다. 또한, 우리는 training time에 meta-learning을 통해 test time에서의 학습을 위한 model의 initialization을 개선합니다.

전반적으로, Test-Time Training(TTT)의 한 형태인 우리의 방식은 이전 형태들과 대조적으로 test time(next-token prediction을 통해)과 training time(meta-learning을 통해) 모두에서 End-to-End(E2E)입니다.

우리는 scaling 속성에 중점을 두고 광범위한 실험을 수행합니다. 특히 164B tokens로 학습된 3B models의 경우, 우리의 방식(TTT-E2E)은 Transformer with full attention과 동일한 방식으로 context length에 따라 scales되는 반면, Mamba 2 및 Gated DeltaNet과 같은 다른 방식들은 그렇지 않습니다.

그러나 RNNs와 유사하게, TTT-E2E는 context length에 관계없이 constant inference latency를 가지므로, 128K context에서 full attention보다 $2.7\times$ 더 빠릅니다. 우리의 code는 공개되어 있습니다.

 

 

Abstract 핵심 정리 노트 (Target: AI Researchers)

이 논문은 Architecture의 변형이 아닌 Learning Paradigm의 전환을 통해 Long-context 문제를 해결하고자 한 연구입니다.

  • Problem Formulation의 전환:
    • Long-context language modeling을 architecture design 문제가 아닌 Continual Learning 문제로 재정의했습니다.
    • 복잡한 구조 대신 standard Transformer with sliding-window attention을 그대로 사용합니다.
  • Core Mechanism: TTT-E2E (Test-Time Training):
    • Test Time: 주어진 context에 대해 next-token prediction을 수행하며 실제로 gradient step을 밟아 weights를 업데이트합니다. 즉, context 정보를 hidden state가 아닌 weights 자체로 압축(compressing)합니다.
    • Training Time: Meta-learning을 도입하여, test time에서의 학습(adaptation)이 더 효율적으로 이루어지도록 model의 initialization을 최적화합니다.
  • Performance & Scaling:
    • Scalability: 3B models, 164B tokens 실험 기준, Mamba 2나 Gated DeltaNet과 달리 Transformer with full attention과 동일한 scaling 경향성을 보입니다.
    • Efficiency: RNNs와 유사하게 context length에 상관없이 constant inference latency를 가집니다.
    • Speed: 128K context 기준으로 full attention 대비 2.7$\times$ 더 빠른 속도를 달성했습니다.

쉬운 설명 : "기억력" 대신 "학습력"을 사용하다

기존의 Long-context 모델들이 정보를 다루는 방식을 '시험 공부'에 비유하자면 다음과 같습니다.

  1. 기존 Transformer (Full Attention): 오픈북 테스트입니다. 질문을 받을 때마다 참고서(Context)의 모든 페이지를 처음부터 끝까지 다시 훑어봅니다. 정확도는 높지만, 책이 두꺼워질수록(Context가 길어질수록) 시간이 엄청나게 오래 걸립니다.
  2. 기존 RNN/State Space Model (Mamba 등): 요약 노트입니다. 책을 읽으면서 중요한 내용을 작은 노트(Hidden state)에 요약해 둡니다. 빠르긴 하지만, 노트 공간이 부족하면 이전 내용을 잊어버리거나(Long-term dependency issue) 정보가 손실됩니다.
  3. 이 논문의 방식 (TTT): **"벼락치기 암기"**입니다.
    • 이 모델은 책을 읽으면서 내용을 노트에 적는 게 아니라, 아예 뇌의 시냅스(Weights)를 재배열해서 내용을 머릿속에 '체화'시켜 버립니다.
    • Test time에 실시간으로 모델 자체가 '학습(Training)'을 하기 때문에, 긴 내용을 다 읽고 나면 그 내용은 이미 모델의 일부가 되어 있습니다.
    • 덕분에 질문이 들어오면 책을 다시 뒤질 필요 없이(Full Attention 불필요), 즉각적으로 대답할 수 있어 속도가 매우 빠릅니다(Constant latency). 또한, 단순 요약보다 훨씬 정교하게 정보를 기억합니다.

 

 

더보기

igure 1 Analysis: Scaling with Context Length

이 차트는 Context length(입력 길이)가 8K에서 128K로 증가함에 따라 각 model의 성능(Left)과 속도(Right)가 어떻게 변하는지를 비교합니다.

모든 model은 3B parameters를 가지며 164B tokens로 학습되었습니다.

Left: Test Loss Scaling (Performance)

  • Y-axis 의미: Loss $\Delta$ with full attention. 즉, Transformer with full attention의 loss를 0으로 기준 잡았을 때의 상대적인 성능 차이입니다. 값이 낮을수록(아래에 있을수록) 성능이 좋으며, 0에 가까우면 full attention과 대등한 성능임을 의미합니다.
  • Analysis:
    • Transformer with full attention (Orange): 기준선이므로 $y=0$에서 평탄하게 유지됩니다. 긴 context에서도 가장 이상적인 성능을 보입니다.
    • Sliding-Window Attention / SWA (Green): Context가 길어질수록 정보 손실이 커져 loss가 급격히 증가(성능 하락)합니다. TTT-E2E의 base architecture임에도 불구하고 가장 나쁜 성능을 보입니다.
    • Mamba 2 & Others (Purple, Brown): Context가 길어짐에 따라 full attention과의 격차가 벌어지며 성능이 저하되는 경향을 보입니다.
    • TTT-E2E (Blue, Ours): 유일하게 Transformer with full attention과 평행한 라인을 유지합니다. 심지어 $y < 0$ 영역에 위치하여, 긴 context에서도 full attention보다 미세하게 더 나은(혹은 대등한) loss를 기록하고 있습니다. 이는 TTT 방식이 SWA의 한계를 완전히 극복했음을 보여줍니다.

Right: Latency Scaling (Efficiency)

  • Y-axis 의미: Latency (sec per 1k tokens). 즉, inference 시 처리 속도를 나타냅니다. 낮을수록 빠릅니다.
  • Analysis:
    • Transformer with full attention (Orange): Context length가 길어질수록 연산량이 급증하여 latency가 가파르게 상승합니다. 128K에서는 매우 느립니다.
    • RNN-like models (Green, Purple, Brown): Mamba 2나 SWA 등은 context 길이에 상관없이 일정한 속도(constant latency)를 유지합니다.
    • TTT-E2E (Blue, Ours): 초기 overhead로 인해 짧은 context에서는 Mamba 2보다 약간 느리지만, 그래프가 **평탄(flat)**합니다. 즉, context가 아무리 길어져도 속도가 느려지지 않습니다.
    • Result: 128K context 기준으로, full attention 대비 2.7$\times$ 더 빠른 속도를 달성하면서도(Right), 성능은 full attention과 동등함(Left)을 입증했습니다.

Summary

이 Figure는 TTT-E2E가 **"Transformer의 성능"**과 **"RNN의 속도(Constant Latency)"**라는 두 마리 토끼를 모두 잡았음을 보여줍니다. 특히 최악의 성능을 보이는 SWA(Green lines) architecture를 기반으로 사용했음에도, Test-Time Training을 통해 최고의 성능을 내는 TTT-E2E(Blue lines)로 변모시킨 점이 핵심입니다.

 

1 Introduction

더보기

인간은 정확한 세부 사항에 대한 불완전한 기억에도 불구하고, 평생 동안 더 많은 경험을 통해 스스로를 발전시킬 수 있습니다. 당신의 첫 번째 machine learning 강의를 생각해보십시오: 강의 중 강사의 첫 마디는 기억나지 않을 수 있지만, 그 강의가 수년 전에 있었다 하더라도 당신이 배운 직관은 아마도 이 논문을 이해하는 데 도움을 주고 있을 것입니다.

반면에, self-attention을 갖춘 Transformers는 여전히 수년간의 인간 경험에 상응하는 long context를 효율적으로 처리하는 데 어려움을 겪고 있는데, 이는 부분적으로 그들이 거의 손실 없는 회상을 위해 설계되었기 때문입니다. 전체 context에 대한 self-attention, 즉 full attention은 모든 새로운 token에 대해 이전의 모든 tokens의 keys와 values를 훑어봐야 합니다. 결과적으로, 이는 모든 세부 사항에 쉽게 주의를 기울이지만, token 당 비용은 context length에 따라 선형적으로 증가하며 빠르게 감당하기 힘든 수준이 됩니다.

Transformers에 대한 대안으로, Mamba 2 및 Gated DeltaNet과 같은 RNNs는 token 당 비용이 일정하지만, Figure 1에서 보여지는 것처럼 더 긴 context에서는 효과가 떨어집니다. 일부 최신 architectures는 sliding window로 full attention을 근사하거나, attention과 RNN layers를 함께 쌓기도 합니다. 그러나 이러한 기법들은 더 나은 성능을 달성하기 위해 더 긴 context를 사용하는 데 있어 여전히 full attention보다 language modeling에서 효과가 떨어집니다.

Token 당 비용이 일정하면서도 language modeling을 위한 효과적인 방법을 어떻게 설계할 수 있을까요? 구체적으로, 서두의 예시처럼 모든 세부 사항을 기억하지 않고도 어떻게 더 긴 context에서 더 나은 성능을 달성할 수 있을까요? 핵심 메커니즘은 압축(compression)입니다. 예를 들어, 인간은 방대한 양의 경험을 뇌로 압축하여, 많은 세부 사항은 생략하면서 중요한 정보는 보존합니다. Language models의 경우, next-token prediction을 통한 training 또한 방대한 양의 데이터를 weights로 압축합니다. 그렇다면 주어진 context에 대해 next-token prediction을 통해 test time에도 language model을 계속 training한다면 어떨까요?

Dynamic evaluation이라 알려진 오래된 아이디어와 유사한 이러한 형태의 Test-Time Training (TTT)에는 여전히 빠진 조각이 있습니다: Training time에, 우리는 TTT 이후의 loss가 아니라 있는 그대로(out of the box)의 loss에 대해 model을 최적화하고 있었습니다. 이러한 불일치를 해결하기 위해, 우리는 standard pre-training 대신 meta-learning을 통해 TTT를 위한 model의 initialization을 준비합니다. 구체적으로, 각 training sequence는 먼저 마치 test sequence인 것처럼 취급되어, 우리는 inner loop에서 이에 대해 TTT를 수행합니다. 그런 다음 우리는 많은 독립적인 training sequences에 걸쳐 TTT 이후의 loss를 평균내고, outer loop에서 gradients of gradients를 통해 TTT를 위한 model의 initialization에 대해 이 평균을 최적화합니다.

요약하자면, 우리의 방법은 두 가지 측면에서 end-to-end입니다. 우리의 inner loop는 long-context TTT에 대한 이전 연구들과 대조적으로 network의 끝에서 next-token prediction loss를 직접 최적화합니다; Subsection 2.4는 우리 방법의 대안적 유도를 통해 이러한 차이를 설명합니다. 게다가, 앞서 논의한 바와 같이 dynamic evaluation과는 대조적으로, 우리의 outer loop는 TTT 이후의 final loss를 직접 최적화합니다. 우리의 주요 결과들은 Figure 1에 강조되어 있으며, 나머지는 Section 3에 제시되어 있습니다.

TTT의 개념적 프레임워크는 long context를 넘어선 많은 응용 분야와, meta-learning이 없는 많은 형태들을 포함하여 긴 역사를 가지고 있습니다. 우리의 연구는 또한 fast weights에 관한 문헌, 특히 우리의 high-level 접근 방식을 공유하는 Clark et al.의 연구에서 영감을 받았습니다. Section 4에서는 related work를 자세히 논의합니다.

 

Introduction 핵심 정리 노트 (Target: AI Researchers)

이 섹션은 기존 Long-context model의 한계를 지적하고, 이를 해결하기 위해 Architecture가 아닌 Learning algorithm 관점에서의 새로운 접근법(TTT)을 제안합니다.

  • Motivation & Problem Statement:
    • Human Memory vs. Transformer: 인간은 세부 사항을 완벽히 기억하지 못해도(imperfect recall) 경험을 직관으로 압축하여 활용합니다. 반면, **Transformer(Full Attention)**는 모든 과거 token의 KV를 스캔하는 'Lossless recall' 방식이라 context length에 따라 비용이 선형적으로 증가합니다.
    • Limitations of Alternatives: RNNs(Mamba 2, Gated DeltaNet)는 constant cost를 가지지만 long context 성능이 떨어집니다. Sliding window나 Hybrid 방식도 full attention의 성능에는 미치지 못합니다.
  • Proposed Solution: Compression via Test-Time Training (TTT):
    • Context 정보를 KV cache나 Hidden state에 저장하는 대신, Weights 자체로 압축(Compression)합니다.
    • Key Idea: Test time에 주어진 context에 대해 next-token prediction을 수행하며 model을 계속 학습시킵니다.
  • Methodology: Meta-learning for TTT:
    • Challenge: 일반적인 Pre-training은 TTT 이후의 loss가 아니라, 초기 상태의 loss(out of the box)를 최적화하므로 mismatch가 발생합니다.
    • Solution: Meta-learning을 도입하여 TTT가 잘 수행되도록 Initialization을 최적화합니다.
      • Inner loop: Training sequence를 test data로 가정하고 TTT 수행.
      • Outer loop: TTT 수행 후의 average loss를 구하고, gradients of gradients를 통해 initialization을 업데이트.
  • Distinctions (Novelty):
    • 기존의 Dynamic evaluation과 달리, TTT 이후의 final loss를 직접 최적화하는 Meta-learning 프레임워크를 갖췄습니다.
    • 이전의 Long-context TTT 연구들과 달리, network의 끝단에서 next-token prediction loss를 직접 최적화하는 완전한 End-to-End 방식입니다.

쉬운 설명 : "참고서"를 버리고 "뇌"를 개조하다

이 논문의 서론은 인간의 학습 방식과 기존 AI 모델의 차이점을 비교하며 시작합니다.

  1. 인간의 방식 (압축과 직관): 우리는 몇 년 전에 들은 수업의 모든 단어를 기억하진 못합니다. 하지만 그 수업을 통해 얻은 '직관'은 뇌 속에 남아서 현재의 문제를 푸는 데 도움을 줍니다. 즉, 정보를 뇌(Weights) 속에 압축해서 저장합니다.
  2. 기존 AI (Full Attention)의 방식: 반면 기존 Transformer는 과거의 모든 기록(Context)을 완벽하게 다 들고 다니면서, 매번 다시 찾아봅니다. 기억력은 완벽하지만, 기억할 양이 많아지면 찾아보는 시간이 너무 오래 걸려서 비효율적입니다.
  3. 이 논문의 접근 (TTT):
    • "시험 시간에도 공부를 계속하자": 모델이 새로운 긴 글을 읽을 때, 단순히 읽고 지나가는 게 아니라 그 자리에서 즉석으로 '학습(Training)'을 해서 내용을 모델의 뇌(Weights)에 새겨버립니다.
    • "공부 잘하는 뇌 만들기(Meta-learning)": 단순히 공부만 시키는 게 아니라, **"어떤 새로운 책을 줘도 순식간에 내용을 흡수할 수 있는 상태"**가 되도록 모델의 초기 상태(Initialization)를 훈련 과정에서 미리 만들어 둡니다.

결국 이 방법은 긴 문맥을 외부 메모리에 저장하는 게 아니라, 모델 자체가 실시간으로 진화하며 정보를 흡수하는 방식입니다.

 

 

2 Method

더보기

Next-token prediction의 표준 작업을 고려해보십시오. 이는 test time에 두 단계로 구성됩니다:

• Prefill: tokens $x_0, x_1, ..., x_T$가 주어졌을 때 $T+1$을 conditioning함, 여기서 $x_0$는 Beginning of Sequence () token입니다.

• Decode: next token의 모든 가능한 instantiations에 대한 distribution $\hat{p}_{T+1}$을 예측함.

Test loss는 $CE(\hat{p}{T+1}, x{T+1})$이며, 여기서 CE는 cross entropy이고 $x_{T+1}$은 자연적으로 생성됩니다. 설명의 편의를 위해, 우리는 먼저 $T+1$ tokens를 prefilling하고 single token을 decoding하는 작업에 초점을 맞춥니다. 이 설정에서, full context에 대한 self-attention, 즉 full attention은 prefill에 $O(T^2)$, decode에 $O(T)$의 computational complexity를 가집니다. 이제 우리는 prefill에 $O(T)$, decode에 $O(1)$을 갖는 Test-Time Training (TTT)을 사용하는 우리의 방법에 대해 논의합니다.

2.1 TTT via Next-Token Prediction

우리의 주요 방법에 동기를 부여하기 위해, Subsection 2.3의 중간까지 발전시킬 toy example을 소개합니다. 이 toy example은 다소 어리석은 architecture에 기반합니다: 모든 self-attention layers가 제거되고 오직 MLP layers만 남은 Transformer입니다. 우리의 toy baseline – 이 architecture를 language modeling에 태연하게 적용한 것 – 은 이전 tokens에 대한 기억이 없기 때문에 효과적으로 bigram입니다. 우리의 목표는 다른 sequence modeling 구성 요소들의 교란 없이 TTT의 효과를 고립시켜 이해하는 것입니다.

우리의 baseline architecture에 약간의 기억을 주는 한 가지 방법은 context에 대해 그것을 훈련시키는 것입니다. Standard pre-training과 유사하게, 우리는 연습으로서 매 $t=1, ..., T$마다 $\hat{p}_t$를 예측하고 $x_t$와 비교할 수 있습니다. 구체적으로, baseline architecture를 weights $W$를 가진 $f$로 표기하면, 시간 $t$에서의 standard next-token prediction loss는 다음과 같이 쓸 수 있습니다:

$$\ell_t(W) = \text{CE}(f(x_{t-1};W), x_t). \quad (1)$$

우리는 매 $t=1, ..., T$마다 순차적으로 gradient descent를 사용하여 test time에 $W$를 업데이트합니다:

$$W_t = W_{t-1} - \eta\nabla\ell_t(W_{t-1}), \quad (2)$$

여기서 $\eta$는 learning rate이고, $W_0$는 test time에서의 initial weights입니다. 마지막에, 우리는 단순히 $\hat{p}_{T+1} = f(x_T; W_T)$를 출력합니다. 우리는 Figure 2의 왼쪽 패널에 이러한 형태의 TTT를 묘사했습니다: 우리의 toy baseline은 위쪽 화살표만 사용하는 반면, TTT는 뒤쪽 및 수평 화살표를 추가합니다.

2.2 Learning to (Learn at Test Time)

이제 우리는 동일한 toy example의 범위 내에서, (training-time) training 후이지만 test-time training 전인 initial weights $W_0$가 어떻게 얻어지는지 고려합니다. 정의에 따르면, 우리의 test-time training loss $\ell_t(W_{t-1})$은 $x_0, ..., x_{t-1}$을 conditioning하고 $x_t$를 예측하려는 next-token prediction task에 대한 test loss이기도 합니다. 따라서 sequence $X = (x_1, ..., x_T)$에 대한 test loss는 다음과 같습니다:

$$L(W_0;X) = \frac{1}{T}\sum_{t=1}^{T} \ell_t(W_{t-1}) = \frac{1}{T}\sum_{t=1}^{T} \text{CE}(f(x_{t-1};W_{t-1}), x_t). \quad (3)$$

Test time에 낮은 $L(W_0;X)$를 생성할 $W_0$를 얻기 위해, 가장 직접적인 접근 방식은 대규모 training set of sequences에 대해 평균적으로 training time에 동일한 loss를 최적화하는 것입니다. 이 직접적인 접근 방식은 training loss가 test loss와 일치하는 End-to-End (E2E) training의 예시입니다. TTT가 이런 방식으로 훈련된 $W_0$를 사용할 때, 우리는 이를 TTT-E2E라고 부릅니다.

대조적인 예로, $W_0$가 test time에 업데이트될 것이라는 점을 고려하지 않고 static model의 training loss를 순진하게 모방하는 또 다른 접근 방식을 고려해 보십시오:

$$L_{\text{naive}}(W_0;X) = \frac{1}{T}\sum_{t=1}^{T} \ell_t(W_0). \quad (4)$$

이 접근 방식은 E2E가 아닙니다. 왜냐하면 model의 training 시와 test 시의 행동 사이에 불일치가 있기 때문입니다. 결과적으로, 우리는 $L_{\text{naive}}$의 minimizer가 또한 낮은 test loss $L$을 생성할 것이라는 보장을 거의 할 수 없습니다. 우리는 이 접근 방식을 TTT-naive라고 부릅니다. 이는 dynamic evaluation 문헌에서 주류 접근 방식이었습니다.

Figure 2의 오른쪽 패널은 $t=1, ..., 128$에 대해 많은 test sequences에 걸쳐 평균낸 token-level test loss $\ell_t$를 도표로 나타냅니다. 지금까지 우리는 네 가지 방법을 논의했습니다: Transformer with full attention (주황색), attention이 없는 우리의 toy baseline (초록색), TTT-naive (회색), 그리고 TTT-E2E ($b=1$인 경우 파란색; $b=16$인 변형은 Subsection 2.3에서 다룰 것입니다); 실험 설정의 세부 사항은 Appendix A를 참조하십시오. TTT-naive는 toy baseline보다 아주 약간 더 나은 성능을 보이는 반면, TTT-E2E는 거의 full attention만큼 좋은 성능을 보입니다. 특히, 시간이 지남에 따라 test loss가 감소하는 것에서 알 수 있듯이, TTT-E2E는 next token을 더 잘 예측하기 위해 더 많은 context를 효과적으로 사용할 수 있습니다.

Gradient-based optimization을 위해, E2E $L$에 대한 $\nabla L(W_0)$를 계산하는 것은 Equation 2의 update rule 자체가 gradient 연산을 포함하기 때문에 gradients of gradients를 계산하는 것을 수반합니다. 다행히도, automatic differentiation을 위한 현대적인 프레임워크들은 최소한의 오버헤드로 gradients of gradients를 효율적으로 계산할 수 있습니다. 일단 $\nabla L(W_0)$가 계산되면, 우리는 그것을 standard optimizers에 연결할 수 있습니다. Meta-learning 분야에서, $L$에 대한 gradient steps는 outer loop라고 불리며, $\ell$에 대한 것은 inner loop라고 불립니다.

현재 버전의 TTT-E2E는 long context에서 large models에 대해 여전히 두 가지 문제를 가지고 있습니다. 첫 번째는 efficiency인데, 우리의 inner loop가 병렬화될 수 없는 많은 단계들을 가지기 때문입니다. 두 번째는 stability인데, inner loop의 각 gradient step이 오직 single token에만 의존하여 우연에 의한 gradient explosion으로 쉽게 이어질 수 있기 때문입니다. 다음 subsection에서는 이 두 가지 문제를 해결합니다.

2.3 Mini-Batch TTT and Sliding Window

위의 두 문제는 공통된 원인을 공유합니다: Equation 2는 mini-batch 대신 online gradient descent를 수행합니다. 크기 $T$의 training set이 주어졌을 때, 표준 관행은 그것을 크기 $b$의 $T/b$ batches로 분할하고(나누어 떨어짐을 가정), batch 당 하나의 gradient step을 밟는 것입니다. $b=1$인 online gradient descent와 비교하여, 더 큰 $b$는 parallelism과 stability 모두를 개선하는 것으로 알려져 있습니다. 우리는 동일한 이점을 위해 mini-batch 아이디어를 TTT에 적용할 수 있습니다. $x_1, ..., x_T$를 포함하는 (test-time) training set이 주어졌을 때, 우리는 Equation 2를 다음과 같이 일반화합니다:

$$W_i = W_{i-1} - \eta \frac{1}{b} \sum_{t=(i-1)b+1}^{ib} \nabla\ell_t(W_{i-1}), \quad (5)$$

여기서 $i=1, ..., T/b$ (나누어 떨어짐을 가정)이며, 그 후 $\hat{p}{T+1} = f(x_T; W{T/b})$를 출력합니다. 또한, (training-time) training이 test-time training에서의 변화를 반영하도록, 우리는 Equation 3 또한 다음과 같이 일반화합니다:

$$L(W_0;X) = \frac{1}{T} \sum_{i=1}^{T/b} \sum_{t=(i-1)b+1}^{ib} \ell_t(W_{i-1}). \quad (6)$$

$b=1$은 Equation 2와 3을 복원한다는 점에 유의하십시오.

그러나, mini-batch TTT를 갖춘 우리의 model은 Figure 2에서 $b=16$인 빨간색 선으로 예시된 바와 같이 각 batch 내에서 다시 bigram이 됩니다. 예를 들어, $x_1, ..., x_b$를 포함하는 첫 번째 mini-batch를 고려하십시오. 모든 예측 $\hat{p}t = f(x{t-1};W_0)$는 $W_{t-1}$ 대신 $W_0$로 만들어지므로, 우리는 $\hat{p}_t$가 더 많은 context, 즉 $t-1$까지의 모든 tokens를 놓침에 따라 $\ell_t(W_0)$가 $t$와 함께 증가하는 것을 관찰합니다. 이 관찰은 context를 놓치지 않는 예측이 mini-batch 내의 첫 번째와 두 번째 예측뿐인 모든 mini-batch 내에서 유효합니다. 이러한 증가하는 losses는 TTT에 대해 더 나쁜 gradient steps를 생성하며, 이는 궁극적으로 파란색 선과 비교하여 보라색 선의 더 나쁜 성능으로 이어집니다.

이 문제를 해결하기 위해, 우리는 마침내 toy example을 넘어 우리의 architecture를 sliding-window attention layers로 증강합니다. 우리의 toy example은 self-attention layers를 완전히 제거한 반면, 우리의 main method는 그것들을 고정된 window size $k$로 제한할 뿐입니다. $T=128K$인 우리의 주요 결과를 위해, 우리는 window size $k$를 $8K$로, TTT mini-batch size $b$를 $1K$로 설정합니다. TTT가 weights를 업데이트할 기회를 갖기 전에 model이 각 mini-batch 내의 context를 기억할 수 있도록 $k \ge b$로 설정하는 것이 중요합니다.

Baseline architecture의 이러한 수정은 우리의 main method를 완성합니다. 다음으로, 우리는 세 가지 구현 세부 사항을 소개하고, multiple tokens를 decoding하는 작업을 고려합니다. 구현 세부 사항으로 완성된 우리의 main method는 Figure 3의 왼쪽 패널에 묘사되어 있습니다.

2.3.1 Implementation Details

우리가 보고한 결과를 달성하기 위해서는 세 가지 구현 세부 사항이 필수적입니다. 우리는 Section 3의 ablations를 통해 이 세부 사항들을 정당화할 것입니다. 그러나 이것들이 단지 우리 실험 설정의 인공물(artifacts)일 뿐이며, 다른 설정에서는 다른 설계 선택이 더 적합할 수도 있습니다.

TTT only the MLP layers. 현대의 Transformers는 반복되는 blocks로 구성되어 있으며, 각 block은 full attention layer (우리는 이를 sliding window attention으로 교체했습니다), MLP layer, 그리고 몇 개의 normalization layers로 구성됩니다. 우리는 TTT 동안 embedding layers, normalization layers, attention layers를 동결(freeze)하는데, inner loop에서 이들을 업데이트하는 것이 outer loop의 불안정성을 야기하기 때문입니다. 따라서 TTT 중에는 오직 MLP layers만이 업데이트됩니다.

TTT only 1/4 of the blocks. 일반적으로 저장 공간이 클수록 압축 과정에서 손실되는 정보가 적습니다. 우리의 경우 정보는 context이고, 저장 공간은 업데이트된 MLP layers입니다. 그러나 더 많은 layers를 업데이트한다는 것은 gradients를 back-propagate하기 위해 더 많은 연산이 필요함을 의미합니다. 따라서 우리는 Section 3의 ablations에서 설명할 것처럼 계산 비용과 context length에 따른 확장 능력(scalability) 사이의 직관적인 trade-off를 가집니다. Ablations에 따라 우리는 마지막 1/4의 blocks만 TTT하기로 선택했지만, 특히 훨씬 더 긴 context를 가진 다른 실험 설정에서는 다른 선택이 필요할 수 있습니다.

Two MLP layers per block. TTT의 우려 사항 중 하나는 pre-training 중에 학습된 지식을 잊어버리는 것입니다. 우리는 이 우려를 해결하기 위해 가장 간단한 방법을 채택합니다. TTT 동안 업데이트되는 blocks에서, 우리는 pre-trained knowledge를 위한 "안전한" 저장소로서 static, second MLP layer를 추가합니다. Baselines와의 공정한 비교를 위해, 우리는 (TTT 동안 동결되는 것들을 포함하여) 전체 네트워크에 걸쳐 MLPs의 hidden dimension을 줄여서 총 parameters 수가 동일하게 유지되도록 합니다.

2.3.2 Decoding Multiple Tokens

지금까지 우리는 $T+1$ tokens (를 $x_0$로 포함)를 prefilling하고 single token을 decoding하는 작업에 집중했습니다. 이제 우리는 multiple tokens를 decoding하는 것을 고려하는데, 이에 대해 우리의 방법은 자연스러운 확장을 허용합니다: Decoded tokens가 TTT mini-batch를 완전히 채웠을 때만 gradient step을 밟습니다. 예를 들어, $T$가 $b$로 나누어 떨어진다고 가정하면, TTT는 정확히 $T/b$ mini-batches 만에 prefilled tokens를 소진합니다. 그러면 우리의 방법은 다음 $b$ tokens를 decoding할 때 특별한 것을 할 필요가 없습니다. 그 후, 이 decoded tokens 배치에 대해 TTT를 수행하고, 업데이트된 weights를 사용하여 decoding을 계속합니다.

2.4 Alternative Derivation

이 subsection은 Key-Value Binding (KVB)에 기반한 long-context TTT에 대한 이전 연구에서 시작하여 우리 main method의 대안적 유도를 논의합니다. 핵심 단계는 그들의 layer-wise reconstruction loss를 standard next-token prediction loss로 대체하여, TTT가 test time에 E2E가 되도록 하는 것입니다. 이 유도는 Section 3의 결과를 이해하는 데 필요하지 않지만, 우리 방법이 RNNs에 관한 문헌과 어떻게 연결되는지에 대한 추가적인 통찰력을 제공합니다.

2.4.1 Starting Point: Key-Value Binding

Context를 model의 weights로 압축한다는 동일한 아이디어에 기반하여, 이전 연구는 self-attention을 대체할 수 있는 sequence modeling layer를 구성하기 위해 TTT를 사용합니다. Self-attention이 모든 이전 token의 key와 value를 cache에 명시적으로 저장하여 연관시키는 반면, 이전 연구는 test time에 각 value를 key로부터 예측하도록 학습함으로써 이러한 연관성을 model 내에 암시적으로 저장할 것을 제안합니다. 나중에 KV Binding으로 알려진 이 학습 목표(learning objective)는 MesaNet, Titans, Nested Learning과 같은 많은 인기 있는 TTT 변형들의 핵심 구성 요소였으며; linear attention과 Gated DeltaNet과 같은 많은 변형들도 이 관점에서 유도될 수 있습니다.

구체적으로, layer $l$에서의 input embeddings $x^{(l)}_t$가 주어졌을 때, TTT-KVB의 기본 형태는 각 $t=1, ..., T$에서 다음 loss에 대해 gradient step을 밟습니다:

$$\ell^{(l)}_t(W^{(l)}_{t-1}) = ||g(\theta^{(l)}_K x^{(l)}_t; W^{(l)}_{t-1}) - \theta^{(l)}_V x^{(l)}_t||^2, \quad (7)$$

여기서 $g$는 보통 MLP이고, $W^{(l)}_{t-1}$은 이전 timestep 후의 $g$의 weights이며, $\theta^{(l)}_K$와 $\theta^{(l)}_V$는 Transformers의 key-value projection matrices와 유사한 outer-loop parameters입니다. Gradient step 후, $g$는 업데이트된 weights를 사용하여 output embedding을 생성합니다:

$$z^{(l)}_t = g(\theta^{(l)}_Q x^{(l)}_t; W^{(l)}_t), \quad (8)$$

여기서 $\theta^{(l)}_Q$ 또한 outer-loop parameters의 집합입니다. 위의 메커니즘은 TTT layer로 알려져 있습니다. Network 내부에서 사용될 때, 모든 TTT layer는 자신만의 loss와 weights를 가진 독립적인 단위입니다. Training time에, TTT-KVB의 outer loop는 Equation 6의 TTT-E2E의 것과 동일하며, 모든 TTT layers의 $\theta_K, \theta_V, \theta_Q$를 포함한 모든 outer-loop parameters가 함께 최적화됩니다. Subsection 2.3의 TTT-E2E와 유사하게, TTT(-KVB) layers는 sliding-window attention layers가 선행될 때 (inner-loop) mini-batch gradient descent를 효과적으로 사용할 수 있습니다. 이 hybrid architecture가 우리 유도의 시작점 역할을 합니다.

2.4.2 First Step: Simplified Output Rule

먼저, 우리는 Equation 8의 output rule이 다음과 같이 단순화될 수 있음을 관찰합니다:

$$z^{(l)}_t = g(\theta^{(l)}_K x^{(l)}_t; W^{(l)}_{t-1}), \quad (9)$$

Table 1에서 보여지는 것처럼 실질적인 해가 없습니다. 이 새로운 output rule은 업데이트된 weights와 별도의 입력 $\theta_Q x_t$로 $g$를 다시 호출하는 대신 Equation 7의 $g$의 예측을 output embedding으로 재사용하기 때문에 단순화입니다. 직관적으로, sliding-window attention이 이미 충분한 local context를 제공한다면 업데이트된 weights로 $g$를 호출하는 것은 불필요할 수 있으며, 이전 연구는 $\theta_K$와 $\theta_Q$ 사이의 분리가 또한 불필요할 수 있다고 주장했습니다.

Figure 3의 오른쪽 패널에서, 우리는 이 단순화 후의 TTT-KVB를 묘사합니다. 왼쪽 패널의 TTT-E2E와 비교하여 네 가지 차이점이 있으며, 후자의 두 가지는 구현 세부 사항으로 간주됩니다:

  1. TTT-KVB의 각 Transformer block은 reconstruction loss $\ell^{(l)}$을 갖는 반면, TTT-E2E는 전체 network의 끝에 단일 next-token prediction loss $\ell$을 갖습니다.
  2. TTT-KVB의 각 Transformer block은 추가적인 outer-loop parameters $\theta_K$와 $\theta_V$를 갖습니다.
  3. TTT-KVB는 모든 Transformer block의 MLP layer를 업데이트하는 반면, TTT-E2E는 마지막 1/4의 blocks의 MLP layer만 업데이트합니다.
  4. 그림에는 나타나지 않았지만, TTT-KVB의 업데이트된 MLPs는 self-attention과 동일한 방식으로 multiple heads로 나뉘므로, 이 MLPs는 TTT-E2E의 일반적인 것보다 훨씬 작습니다. 게다가 이 MLPs는 LoRA로 업데이트되므로 유효 용량(effective capacity)은 훨씬 더 작습니다.

그러나 Figure 3은 또한 TTT의 이 두 가지 형태 사이의 유사성을 강조합니다: TTT-E2E와 유사하게, TTT-KVB는 전체 network를 훈련시키는 관점에서 이해될 수 있습니다. 먼저, 모든 위쪽 화살표로 묘사된 것처럼 전체 network를 통한 forward pass가 있습니다. 그 다음 Deeply Supervised Nets 방식의 많은 losses로부터의 기여를 가진 backward pass가 있습니다. 그러나 각 block의 단일 아래쪽 화살표로 묘사된 것처럼, gradients는 오직 하나의 MLP layer에 도달한 후 멈춥니다.

2.4.3 Key Step: E2E at Test Time

이 유도의 핵심 단계는 KVB loss를 next-token prediction loss로 대체하는 것인데, 이는 layer-wise reconstruction losses 없이는 $\theta_K$나 $\theta_V$의 용도 또한 없기 때문에 차이점 1과 2를 함께 제거하는 것을 의미합니다. 이 단계는 우리를 TTT-E2E all layers MH라는 중간 방법으로 이끄는데, 여기서 MH는 multi-head를 의미합니다. Table 1에서 보여지듯이, loss를 대체하는 것은 language modeling 성능을 상당히 향상시켜 본질적으로 우리의 최종 방법 수준에 도달합니다.

이 중간 방법은 이제 test time에 E2E인데, 그 (test-time) training loss가 정확히 token-level test loss $\ell_t$이기 때문입니다. 이 시점에서 우리의 두 유도 사이의 이중성을 인식하는 것은 특히 흥미롭습니다. 우리의 주된 유도는 test time에 E2E인 next-token prediction을 통한 TTT에서 시작하여, Subsection 2.2에서 meta-learning을 통해 training time에 E2E로 만드는 데 집중했습니다. 반면 우리의 대안적 유도는 training time에 E2E인 TTT-KVB에서 시작하여, next-token prediction을 통해 test time에 E2E로 만드는 데 집중했습니다.

2.4.4 Final Step: Larger State with Less Compute

TTT(-KVB) layers는 종종 RNN layers의 한 종류로 간주되며, TTT-KVB는 종종 RNN으로 간주됩니다. 마찬가지로, TTT-E2E 또한 전체 network가 한 번의 backward pass로 업데이트되기 때문에 오직 하나의 RNN layer를 가진다는 점을 제외하면 RNN으로 볼 수 있습니다. RNN의 세 가지 구성 요소 중, 우리는 이 유도의 첫 번째 단계에서 output rule을 수정했고, 두 번째 (핵심) 단계에서 update rule을 수정했습니다. 이제 우리는 hidden state를 수정합니다.

RNN의 long-context 성능을 종종 향상시키는 중요한 요소는 더 큰 hidden state이며, 이는 결과적으로 업데이트하기 위해 더 많은 연산을 필요로 합니다. 우리의 중간 방법인 TTT-E2E all layers MH를 고려하십시오. 만약 우리가 마지막 1/4의 blocks만 업데이트함으로써 차이점 3을 제거한다면, 우리는 test time에 연산을 절약하지만 더 작은 state로 끝나게 됩니다. 그리고 만약 우리가 (LoRA를 갖춘 multi-head MLPs 대신) regular MLPs로 되돌아감으로써 차이점 4를 제거한다면, 우리는 더 많은 연산(및 메모리)을 비용으로 더 큰 유효 state(effective state)를 갖게 됩니다.

그러나 E2E loss를 사용할 때, 이 두 차이점의 trade-offs는 다소 불균형합니다: Block 내의 작은 multi-head MLP를 업데이트하기 위해서는 backward pass가 더 진행되기 위해 그 아래의 attention layer는 말할 것도 없고, 그 위의 큰 MLP를 통해 gradients가 back-propagate되어야 합니다. Upstream gradients를 준비하는 막대한 비용을 감안할 때, 각각 더 큰 hidden state를 포함하는 더 적은 수의 blocks를 업데이트하는 것이 비용 효율적일 것입니다. 실제로, 두 차이점을 모두 제거한 우리의 최종 방법(TTT-E2E)은 TTT-E2E all layers MH와 비교하여 5배 더 큰 hidden state (760M model의 경우 88M vs. 18M)와 절반의 inference latency (H100에서 prefill 시 1K tokens 당 0.0086 vs. 0.017초)를 가집니다.

이 더 큰 state가 실제로 성능을 향상시킬까요? Table 1에서는 이 실험들이 8K context length에서만 수행되었기 때문에 차이를 보기 어렵습니다. Subsection 3.2에서, 우리는 업데이트되는 layers의 수를 제거(ablating)함으로써 context length에 따른 scaling 측면에서 state size의 효과를 조사할 것입니다. 이 ablation은 더 작은 state가 더 나쁜 context scaling으로 이어진다는 것을 명확히 보여줄 것입니다.

 

 

 

 

 

 

 

 

 

 

더보기

Figure 2: Toy Example (TTT의 기본 원리 증명)

이 그림은 복잡한 Transformer 구조 없이, TTT(Test-Time Training)만으로도 모델이 "기억(Memory)"을 가질 수 있다는 것을 보여주는 실험입니다.

Left: Mechanism (작동 원리)

  • Architecture: Self-attention이 제거된 Transformer입니다. 즉, $x_1$을 보고 $x_2$를 예측할 때, 일반적인 경우라면 이전 정보($x_1$)를 참고할 방법이 없어 Bigram (직전 단어만 보고 예측) 수준에 머무릅니다.
  • TTT Process:
    1. Forward (위쪽 화살표): $x_1$이 입력되어 $x_2$를 예측($\hat{p}_2$)합니다.
    2. Loss Calculation: 실제 $x_2$와 예측값 $\hat{p}_2$ 사이의 Loss $\ell_2$를 계산합니다.
    3. Backward & Update (아래쪽 & 수평 화살표): 이 Loss를 기반으로 Gradient step을 밟아 MLP의 Weights를 업데이트합니다 (파란색 박스로 변함).
    4. Result: 업데이트된 Weights는 $x_1$에 대한 정보를 "압축"하여 담고 있습니다. 이제 $x_3$를 예측할 때, 모델은 업데이트된 파란색 MLP를 사용하므로 $x_1$의 문맥을 반영할 수 있게 됩니다.

Right: Performance (성능 비교)

  • Orange Line (TF. with attn): Full Attention을 사용하는 일반적인 Transformer입니다. Loss가 가장 낮고 안정적입니다 (Ideal Baseline).
  • Green Line (TF. w/o attn): Attention도 없고 TTT도 없는 경우입니다. 문맥을 전혀 기억하지 못하므로 Loss가 높고 줄어들지 않습니다 (Bigram Baseline).
  • Blue Line (TTT-E2E $b=1$): 우리의 Toy model에 TTT를 적용한 결과입니다. 놀랍게도 **Full Attention(Orange)**에 근접하는 성능을 보여줍니다. 이는 Weights Update만으로도 문맥 정보를 효과적으로 저장할 수 있음을 증명합니다.
  • Red Line (TTT-E2E $b=16$): Mini-batch TTT를 적용했을 때의 문제입니다. Batch size 16 동안은 Weights가 업데이트되지 않으므로, 그 구간 안에서는 Green Line처럼 문맥을 잊어버려 Loss가 치솟습니다(톱니바퀴 모양).
    • Insight: 이 문제를 해결하기 위해 저자들은 Sliding Window Attention을 도입해야 한다는 결론을 도출합니다 (Figure 3로 연결됨).

Section 2.0 - 2.2 핵심 정리 노트 (Target: AI Researchers)

이 섹션은 TTT(Test-Time Training)의 기본 메커니즘과 이를 위한 Meta-learning 기반의 최적화 방법론(TTT-E2E)을 정의합니다.

  • Computational Complexity Advantage
    • Full Attention: Prefill $O(T^2)$, Decode $O(T)$의 비용 발생.
    • TTT Method: Prefill $O(T)$, Decode $O(1)$로 RNN과 유사한 효율성 달성.
  • 2.1 TTT Mechanism (Toy Example Verification)
    • Hypothesis: Self-attention 없이 MLP layer만 있는 Transformer(사실상 Bigram)라도, Test time에 Context에 대해 학습하면 "기억(Memory)"을 가질 수 있는가?.
    • Method: Test time에 $t=1, ..., T$ 순서대로 Gradient Descent를 수행하여 Weights $W$를 즉시 업데이트.
      • Update Rule: $W_t = W_{t-1} - \eta\nabla\ell_t(W_{t-1})$.
    • Result: Attention이 없는 단순 MLP 모델도 TTT를 적용하면 Full Attention에 근접한 Loss 감소를 보임 (Figure 2 참조).
  • 2.2 TTT-E2E (Meta-Learning for Initialization)
    • Problem of TTT-naive: 기존 Dynamic evaluation 방식은 $W_0$를 Static model처럼 학습($\ell_t(W_0)$ 최소화)시킨 후, Test time에만 업데이트함. 이는 Training과 Test 동작 간의 Mismatch를 유발.
    • Proposed TTT-E2E: Training time부터 "Test time에 학습할 것"을 가정하고 $W_0$를 최적화 (Meta-learning).
      • Objective: 단순한 Prediction loss가 아닌, **TTT 업데이트가 수행된 후($W_{t-1}$)**의 Loss를 최소화.
      • Loss Function: $L(W_0;X) = \frac{1}{T}\sum_{t=1}^{T} \text{CE}(f(x_{t-1};W_{t-1}), x_t)$.
    • Optimization: Outer loop(Training)에서 Inner loop(Test TTT)의 과정을 미분해야 하므로 Gradients of gradients 계산이 필요함.
  • Current Challenges (Limits of 2.2)
    • Section 2.2까지의 모델은 Online Gradient Descent($b=1$)를 사용하므로 Efficiency(병렬화 불가)와 Stability(Gradient Explosion 위험) 문제가 존재하며, 이는 2.3절의 Mini-batch 도입 배경이 됨.

쉬운 설명 : "공부하는 법"을 가르치는 모델

1. 벼락치기 공부법 (Section 2.1)

일반적인 AI 모델(Transformer)은 이미 학습된 지식(Weights)만 가지고 시험을 봅니다. 하지만 이 논문의 TTT 방식은 시험장에 들어가서도 지문(Context)을 읽으면서 즉석에서 뇌세포(Weights)를 뜯어고치며 공부합니다.

저자들은 극단적인 실험을 했습니다. "기억력(Attention)"이 아예 없는 멍청한 모델(Bigram)을 만들고, 이 모델에게 "지문을 읽으면서 실시간으로 뇌를 개조해라"라고 시켰더니, 놀랍게도 기억력이 좋은 모델만큼 똑똑해졌습니다. 즉, "기억"을 "어딘가에 적어두는 것(Cache)"이 아니라 "내면화(Weights Update)"하는 것으로 대체한 것입니다.

2. 실전형 조기 교육 (Section 2.2)

그런데 문제가 있습니다. 모델을 처음 만들 때(Pre-training), "너는 나중에 시험장에서 벼락치기를 하게 될 거야"라고 알려주지 않고 그냥 공부시키면(TTT-naive), 막상 시험장에서 벼락치기를 잘 못할 수 있습니다.

그래서 TTT-E2E 방식은 훈련 단계에서부터 **"벼락치기를 잘하는 뇌 구조(Initialization $W_0$)"**를 만들어줍니다.

  • 기존 방식: "정답을 잘 맞추는 상태"를 만듦.
  • 이 논문의 방식: "새로운 정보를 봤을 때, 순식간에 학습해서 내 것으로 만드는 상태"를 만듦.

이것을 수학적으로 구현하기 위해 **"학습하는 과정을 학습(Learning to Learn)"**하는 메타 러닝 방식을 사용합니다.

 

 

 

 

 

 

 

 

Section 2.3 핵심 정리 노트 (Target: AI Researchers)

이 섹션은 TTT를 실제 Large Scale Model에 적용하기 위해 Efficiency(속도)와 Stability(안정성)를 확보하는 구체적인 방법론을 다룹니다.

  • Mini-Batch TTT Introduction
    • Motivation: $b=1$(Online GD)은 병렬화가 불가능하고, Single token의 노이즈로 인한 Gradient Explosion 위험이 있음.
    • Method: Training set을 size $b$로 나누어 배치 단위로 Update 수행.
      • Update Rule: $W_i = W_{i-1} - \eta \frac{1}{b} \sum_{t=(i-1)b+1}^{ib} \nabla\ell_t(W_{i-1})$
    • New Problem: Batch 내부에서는 Weights가 고정($W_{i-1}$)되므로, Batch 초반부의 정보를 후반부에서 참조하지 못하는 Intra-batch Amnesia 현상 발생 (Figure 2의 Red line, 톱니바퀴 모양 Loss). 사실상 Batch 내에서는 다시 Bigram 모델로 회귀함.
  • Hybrid Architecture: TTT + Sliding Window Attention (SWA)
    • Solution: Sliding Window Attention을 도입하여 "Batch 내부의 단기 기억"을 담당하게 함.
    • Constraint: Window size $k$는 반드시 Batch size $b$보다 커야 함 ($k \ge b$).
      • Role Division: SWA는 Batch 내의 Local Context(Short-term)를 처리하고, TTT는 이를 Weights로 압축하여 Global Context(Long-term)를 처리.
  • Crucial Implementation Details (Ablation-based Choices)
    1. TTT Target: 오직 MLP Layers만 업데이트.
      • Attention, Norm, Embedding은 Freeze. (Inner loop 업데이트 시 Outer loop 불안정성 유발 방지)
    2. Scope: 전체 블록이 아닌 Last 1/4 Blocks만 업데이트.
      • Compute Cost vs. Memory Capacity Trade-off. 적은 수의 블록에 집중하는 것이 Gradient Backprop 비용 대비 효율적임.
    3. Dual MLP Structure: TTT 블록 내에 Static Second MLP 추가.
      • Pre-training 지식의 망각(Catastrophic Forgetting)을 방지하기 위한 "Safe Storage" 역할. (Parameters 수는 타 모델과 동일하게 조정됨)
  • Decoding Strategy
    • Prefill 이후 Decode 단계에서도 생성된 토큰이 Batch $b$를 채울 때마다 Gradient Step을 수행하여 Weights를 지속적으로 업데이트.

쉬운 설명 : "속독"과 "정독"의 조화

1. 한 글자씩 공부하기 vs 문단별로 공부하기 (Mini-Batch)

이전 섹션(2.2)까지의 모델은 글자 하나를 읽을 때마다 뇌(Weights)를 고쳤습니다. 이건 너무 느리고 피곤한 일입니다(비효율적). 그래서 저자들은 **"문단(Batch) 단위로 읽고 나서 한꺼번에 뇌를 고치자"**라고 전략을 바꿉니다.

2. 문단 내에서의 기억상실 문제

그런데 문제가 생깁니다. 한 문단을 읽는 도중에는 뇌를 고치지 않으니, 문단의 첫 번째 문장 내용을 마지막 문장 읽을 때 까먹어버립니다(Batch 내에서는 멍청해짐). 마치 문단을 다 읽기 전까진 이해를 못 하는 것과 같습니다.

3. 단기 기억 장치 추가 (Sliding Window Attention)

이 문제를 해결하기 위해 **"단기 기억 안경(Sliding Window Attention)"**을 씌워줍니다.

  • 안경(SWA): 지금 읽고 있는 문단 내의 내용은 안경을 통해 바로바로 봅니다. (단기 기억)
  • 뇌(TTT): 문단을 다 읽고 나면, 그 핵심 내용을 뇌의 주름(Weights)에 영구적으로 새깁니다. (장기 기억/압축)
  • 조건: 안경이 볼 수 있는 범위($k$)는 최소한 한 문단의 길이($b$)보다는 길어야 끊김 없이 내용을 이어갈 수 있습니다.

4. 뇌 수술의 디테일 (Implementation Details)

뇌를 개조할 때도 규칙이 있습니다.

  • 위험한 부위는 건드리지 마라: 시각 신경(Attention)이나 기본 감각(Embedding)을 건드리면 모델이 미쳐버립니다. 그래서 **추론 능력(MLP)**만 건드립니다.
  • 일부만 개조해라: 뇌 전체를 다 뜯어고치면 너무 대수술이 됩니다. 가장 고차원적인 사고를 하는 뒤쪽 1/4만 수술합니다.
  • 백업용 뇌를 남겨둬라: 너무 새로운 것만 배우다가 원래 알고 있던 상식(Pre-training 지식)을 까먹을 수 있습니다. 그래서 개조하는 부위 옆에 **변하지 않는 보조 뇌(Static MLP)**를 하나 더 붙여서 균형을 맞춥니다.

 

 

 

 

 

 

 

더보기

Figure 3: Computation Graphs (제안하는 모델 구조)

Figure 2의 Toy Example에서 발견한 한계(Batch 내 기억 상실)를 보완하고, 실제 대규모 모델에 적용하기 위한 최종 구조(Left)와 기존 연구(Right)를 비교한 그림입니다.

Left: (a) Our Main Method (TTT-E2E)

이 논문이 제안하는 최종 모델 구조입니다.

  1. Sliding Window Attention (SWA):
    • 모든 블록에 Sliding Window가 적용되어 있습니다. 이는 Figure 2의 Red Line 문제(Batch 내에서 문맥을 놓치는 문제)를 해결해 줍니다. 즉, TTT가 업데이트되기 전의 짧은 문맥은 SWA가 담당합니다.
  2. Selective TTT (Top 1/4 Update):
    • 그림을 보면 아래쪽 $\frac{3L}{4}$ 블록(회색)은 TTT 업데이트를 하지 않고, 위쪽 $\frac{L}{4}$ 블록(파란색)만 업데이트합니다.
    • Disconnected Arrow: 위쪽 블록에서 발생한 Gradient가 아래쪽 블록으로 전파되지 않도록 끊겨 있습니다. 이는 연산 효율성을 높이고(Gradient 계산 줄임), 학습 안정성을 확보하기 위함입니다.
  3. End-to-End (E2E) Loss:
    • 가장 중요한 특징입니다. Loss $\ell_2$가 네트워크의 **가장 끝단(Final Output)**에서 계산됩니다. 즉, "다음 토큰을 잘 맞추는 것"을 목표로 전체 네트워크가 조율됩니다.

Right: (b) Prior Work (TTT-KVB Simplified)

기존의 TTT 연구들(Mamba 등의 기반이 되는 아이디어)과의 차이점입니다.

  1. Layer-wise Loss:
    • Loss $\ell^{(l)}_1$이 각 레이어마다 개별적으로 존재합니다 (Reconstruction loss). 즉, 각 레이어가 독립적으로 입력을 복원하려고 노력합니다.
    • 저자들은 이것이 최종 목표(Next-token prediction)와 직접 연결되지 않는다고 비판합니다.
  2. Complex Parameters:
    • $\theta_K, \theta_V$와 같은 추가적인 파라미터가 필요하며 구조가 더 복잡합니다.

요약

  • Figure 2는 **"Weights Update로 문맥을 기억할 수 있다"**는 TTT의 핵심 가설을 증명합니다.
  • Figure 3는 **"SWA로 단기 기억을 보완하고, 상위 레이어만 E2E로 학습시켜 장기 기억과 효율성을 모두 잡았다"**는 이 논문의 최종 방법론을 보여줍니다.

 

 

Section 2.4 핵심 정리 노트 (Target: AI Researchers)

이 섹션은 TTT-E2E를 기존의 RNN/Linear Attention 계열 연구(특히 TTT-KVB)의 연장선상에서 재해석(Alternative Derivation)하며, State Size와 Compute Efficiency 간의 Trade-off를 최적화하는 과정을 설명합니다.

  • Starting Point: TTT-KVB (Key-Value Binding)
    • Concept: Self-attention의 Explicit KV Cache 대신, Test time에 $K \to V$ 매핑 함수 $g$를 학습하여 Implicit하게 저장.
    • Loss: Layer-wise Reconstruction Loss 사용 ($\ell^{(l)}_t = ||g(K_t) - V_t||^2$).
    • Mechanism: 각 Layer가 독립적인 TTT Unit으로 동작.
  • Step 1: Simplified Output Rule
    • 기존 KVB는 Update 후 별도의 Query projection($\theta_Q$)을 통해 Output을 계산했으나, 이를 **Loss 계산 시 사용한 예측값($g(\theta_K x; W_{t-1})$)**으로 대체.
    • Justification: Sliding Window Attention이 Local Context를 제공하므로, Update 직후의 Weight로 다시 계산하는 연산은 Redundant함 (Table 1 참조).
  • Step 2: From Layer-wise to Global E2E Loss (Critical Shift)
    • Action: Layer별 Reconstruction Loss를 Network 전체의 Next-Token Prediction Loss로 대체.
    • Implication: $\theta_K, \theta_V$와 같은 불필요한 Projection 제거.
    • Duality Discovery:
      • Main Derivation (Sec 2.1~2.3): Test-time E2E(Next-token)에서 시작 $\to$ Meta-learning을 통해 Training-time E2E 달성.
      • Alternative Derivation (Sec 2.4): Training-time E2E(KVB)에서 시작 $\to$ Loss 변경을 통해 Test-time E2E 달성.
  • Step 3: Strategic Resource Allocation (State Size vs. Compute)
    • Comparison:
      • Typical RNN/KVB: 모든 Layer의 Small State(Multi-head LoRA 등)를 업데이트.
      • TTT-E2E: 상위 1/4 Layer의 Large State(Full MLP)만 업데이트.
    • Efficiency Analysis: 하위 Layer의 작은 State를 업데이트하기 위해 상위 Layer 전체를 Backprop하는 것은 비효율적임 (Gradient Computation Overhead).
    • Result: 상위 Layer에 집중함으로써 $5\times$ Larger Hidden State를 확보하면서도 $2\times$ Faster Inference 달성 (H100 기준).

쉬운 설명 : "암기 카드" 버리고 "실전 문제"로 승부하다

1. 기존 방식 (TTT-KVB): 쪽지 시험의 연속

기존의 유사한 연구들은 모델의 각 층(Layer)마다 "암기 카드(Key-Value)"를 주고, "이거 외웠니?"라고 확인하는 쪽지 시험(Layer-wise Loss)을 계속 쳤습니다. 각 층이 자기 할 일만 잘하면 전체적으로 잘할 것이라는 믿음 때문이었죠.

2. 새로운 방식 (TTT-E2E): 수능 문제 풀기

저자들은 이 방식이 비효율적이라고 판단했습니다. 각 층이 암기 카드를 잘 외웠는지 검사하는 과정(Reconstruction Loss)을 없애고, **"그래서 다음 단어가 뭔데?"(Next-Token Prediction)**라는 최종 결과만 따지기로 했습니다.

중간 과정이 어찌 됐든, 정답만 맞히면 되도록 모델 전체를 한 번에 조율하는 것(End-to-End)이 훨씬 강력하다는 것을 발견했습니다.

3. 선택과 집중 (효율성)

또한, 뇌의 모든 부분(All Layers)을 조금씩 고치는 것보다, 가장 고차원적인 사고를 담당하는 부분(Top Layers)을 크고 확실하게(Large MLP) 고치는 것이 낫다는 것을 증명했습니다.

  • 기존: 발가락부터 머리까지 신경 하나하나를 미세하게 조정함. (시간 오래 걸림, 효과 적음)
  • 제안: 뇌의 전두엽만 확실하게 뜯어고침. (시간 절약, 용량 5배 증가, 속도 2배 증가)

 

3 Main Results

더보기

모든 실험은 우리의 public repository에서 제공되는 code와 datasets를 사용하여 재현할 수 있습니다:

https://github.com/test-time-training/e2e.

3.1 Setup

이 논문의 연구적 성격을 고려할 때, 우리의 목표는 large scale에서의 production-level training runs에 정보를 줄 수 있는 small 및 medium scales의 가장 간단한 setups에서 실험하는 것입니다. 일반적으로, 오늘날의 large-scale runs는 보통 두 개 이상의 stages로 구성됩니다:

• 다양한 지식을 포함하는 general dataset에서 short context length로 Pre-training.

• Long sequences의 dataset에서 fine-tuning하여 context length를 확장. 매우 긴 context(예: 1M)에 점진적으로 도달하기 위해, extension은 보통 여러 stages로 더 세분화됩니다.

단순화를 위해, 우리의 training runs는 오직 두 단계로 구성됩니다: 8K context length에서의 pre-training, 그리고 실험에 따라 최대 128K의 final context length에서의 extension fine-tuning.

Datasets. Pre-training을 위해, 우리는 Common Crawl의 엄격하게 필터링된 subset인 DCLM, 구체적으로 DCLM-Baseline을 사용합니다. DCLM-Baseline의 3.8T tokens를 고려하여, 우리는 먼저 우리의 pre-training context length인 8K보다 짧은 모든 documents를 버리고, 남은 것들 중에서 무작위로 sample하여 다양한 크기의 training sets를 구성합니다. 그러나 extension을 위한 우리의 maximum context length인 128K보다 긴 DCLM의 sequences 대부분은 품질이 낮습니다. 그래서 fine-tuning을 위해, 우리는 long-context extension을 위한 standard academic dataset인 Books를 사용합니다. 우리는 또한 language modeling evaluation을 위해 Books의 held-out partition을 사용합니다.

Basic recipe. 우리는 125M에서 3B parameters 범위의 다섯 가지 크기의 models로 실험합니다. 다양한 실험에서의 우리의 configurations와 training hyper-parameters는 모두 Appendix B에 자세히 설명된 단일 basic recipe에서 파생되었습니다. 요약하자면, model configurations와 pre-training을 위한 basic recipe는 GPT-3와 Mamba에서 가져왔습니다; fine-tuning을 위한 basic recipe를 생성하기 위해, 우리는 full attention을 갖춘 Transformer baseline에 대해 grid search를 수행했습니다.

Baselines. 우리는 우리의 방법을 architecture design에서 state-of-the-art approaches를 대표하는 6개의 baselines와 비교합니다. Sliding window를 사용하는 모든 baselines는 동일한 window size $k = 8K$를 사용합니다.

  1. Transformer with full attention: 위에서 논의된 model configurations를 사용함.
  2. Transformer with Sliding-Window Attention (SWA): 모든 full attention layer가 SWA layer로 대체됨. Implementation details가 없는 Subsection 2.3의 우리의 main method 또한 이 architecture에 기반함. Window size $k$는 window size ablations를 제외한 모든 실험에서 8K로 설정됨. Pre-training context length 또한 8K이므로, full attention과 SWA baselines는 extension fine-tuning 전까지 동일함.
  3. Hybrid SWA and full attention (5:1): Gemma 스타일로 5개의 SWA layers 뒤에 1개의 full attention layer가 오는 패턴을 반복함.
  4. Mamba 2: Mamba 2 layers와 SWA layers의 hybrid를 사용하는 인기 있는 RNN; Nemotron-H에서 large scale로 테스트됨.
  5. Gated DeltaNet: Mamba 2와 DeltaNet을 확장하고 Gated DeltaNet layers와 SWA layers의 hybrid를 사용하는 인기 있는 RNN; Kimi Linear에서 large scale로 테스트됨.
  6. TTT-KVB: Key-Value Binding (KVB)을 갖춘 TTT-MLP layers와 SWA layers의 hybrid를 사용하는 인기 있는 RNN; 또한 (simplified output rule이 없는) Subsection 2.4에서의 우리의 시작점임. Titans와 Nested Learning은 유사한 구조를 따름.
  7. 우리는 baselines 1–3을 우리 자신의 방법과 함께 JAX로 구현합니다. Baselines 4–6의 경우, 우리는 저자들이 제공한 official code와 configurations를 사용하며, 가능한 경우 baselines를 개선하기 위해 그들과 상의했습니다. Baselines에 대한 우리의 개선 사항은 Appendix C에서 논의됩니다.

3.2 Ablations on Hyper-Parameters

독자들이 점진적으로 우리 방법에 대한 경험적 직관을 쌓을 수 있도록 돕기 위해, 우리는 가장 간단한 실험들 – Subsection 2.3에서 소개된 hyper-parameters에 대한 ablations로 시작합니다. 모든 ablations에 대해, 우리는 basic recipe를 갖춘 760M model을 사용합니다.

Sliding window size $k$. 이 hyper-parameter는 full attention을 제외한 모든 methods에 존재합니다. 따라서, 우리는 두 가지 대표적인 baselines인 SWA와 Gated DeltaNet에 대해서도 이 ablation을 수행합니다. 놀랍지 않게도, Figure 4의 가장 왼쪽 패널에서 보여지듯이 더 큰 $k$는 세 가지 methods 모두의 성능을 향상시키며, TTT-E2E는 baselines와 비교하여 $k$의 변화에 유사한 sensitivity를 가집니다. 우리는 더 작은 $k$가 runtime을 유의미하게 개선하지 않으므로 $k = 8K$를 default로 선택합니다.

TTT-E2E with full attention. Window size ablation은 Books에 대한 fine-tuning 없이 DCLM에 대한 pre-training만으로 수행되므로, 위의 결과들은 DCLM에서도 평가됩니다. Pre-training context length 또한 8K이므로, $k = 8K$인 SWA는 정확히 full attention이며, $k = 8K$인 TTT-E2E는 full attention 위의 TTT-E2E와 동일하게 됩니다. TTT-E2E가 full attention 위에서도 test loss를 (0.018만큼) 개선할 수 있고, $k$가 증가함에 따라 TTT-E2E와 SWA 사이의 차이가 유의미하게 변하지 않는다는 것을 관찰하는 것은 특히 흥미롭습니다. 이 관찰은 TTT-E2E가 단순히 full attention과 SWA 사이의 차이를 보상하는 것이 아님을 시사합니다; 대신, context length와 같은 다른 요인들이 고정되었을 때 orthogonal improvement를 생성합니다.

TTT mini-batch size $b$. Figure 4의 중간 패널은 1K에서 8K 범위의 TTT mini-batch size $b$로 실험합니다. 이 hyper-parameter는 TTT 관점에서 파생된 methods에 고유하므로, 여기서 의미 있는 비교를 허용하는 유일한 다른 baseline은 TTT-KVB입니다. Window size ablation과 유사하게, models는 pre-training 후 DCLM에서 평가됩니다. TTT-E2E와 TTT-KVB 모두에 대해, 우리는 더 큰 $b$의 선택이 성능을 유의미하게 해친다는 것을 관찰합니다. 그러나, 1K보다 작은 $b$의 선택 또한 우리의 hardware utilization과 stability를 실험하기 어려울 정도로 유의미하게 해칩니다. 따라서, 우리는 $b = 1K$를 default로 선택합니다.

Modified architectures without TTT. $b = 8K$의 선택은 우리의 pre-training context length 또한 8K이므로 TTT를 전혀 하지 않는 것과 동일합니다. 그러나, TTT가 없는 TTT-E2E와 TTT-KVB 둘 다 Figure 3에서 이전에 묘사된 것처럼 Transformer architecture를 약간 수정했기 때문에 full attention을 갖춘 Transformer와는 약간 다릅니다. 그렇다면 이러한 수정들이 TTT 없이도 여전히 중요할까요? Figure 4는 답이 아니오임을 시사합니다. TTT 없이, TTT-E2E(2.825)나 TTT-KVB(2.826)의 loss는 full attention(2.827)과 거의 다르지 않습니다. 이 관찰은 architecture design이 우리의 방법에서 부차적이고 보조적인 역할을 한다는 것을 시사합니다.

3.2.1 Number of Layers Updated

이제 우리는 가장 중요한 ablation으로 넘어갑니다. Subsection 2.3에서 논의된 바와 같이, TTT 동안 업데이트되는 layers의 수는 우리가 context window의 정보를 압축할 수 있는 storage의 양을 제어합니다. 따라서, 우리는 context scaling 측면에서 그 효과를 조사하고, 이 ablation을 Figure 1 (left)의 형식으로 제시합니다. 구체적으로, 각 layer 수에 대해, 우리는 DCLM에서 하나의 checkpoint를 pre-train한 다음 Books에서 다섯 가지 버전을 fine-tune하는데, 각 context length마다 하나씩 하므로 final results는 Books에서 평가됩니다.

우리는 마지막 1/2, 1/4, 1/8의 layers를 업데이트하는 실험을 합니다. 총 24 layers를 가진 우리의 760M model의 경우, 이 비율들은 마지막 12, 6, 3 layers로 해석됩니다. 우리는 또한 오직 final layer만 업데이트하는 실험도 합니다. Figure 4의 가장 오른쪽 패널에서, 우리는 오직 1개 또는 3개의 layers만 업데이트할 때 우리의 방법이 full attention과 같은 방식으로 context length에 따라 scale하지 않음을 관찰합니다. 6개 또는 12개의 layers를 업데이트할 때, 우리의 방법은 scale합니다. 그러나, 12개의 layers를 업데이트하는 것은 6개와 대략 같은 수준의 성능만 냅니다. 따라서, 우리는 model size에 관계없이 항상 마지막 1/4을 업데이트합니다.

3.3 Scaling with Training Compute

일반적으로, training compute에는 두 개의 축이 있습니다: model size와 training tokens의 수입니다. 우리는 full attention 및 Gated DeltaNet과 비교하여 이 축들에 따른 우리 방법의 동작을 조사하고, 결과를 Figure 5에 제시합니다. 우리는 Gated DeltaNet이 고도로 최적화된 training time을 가진 가장 최신 연구이기 때문에 RNN baselines 중 대표로 선택합니다.

Training compute의 효과를 측정하는 하나의 대중적인 관행은 많은 scaling law 논문들에서처럼 pre-training 직후에 pre-training dataset에서 평가하는 것입니다. Figure 5의 왼쪽 패널들에서, 우리는 이 관행을 따르고 pre-training 후 DCLM에서 평가합니다. 그러나 Subsection 3.2에서 논의된 바와 같이, 우리의 window size는 pre-training context length와 같으므로, 우리의 baseline architecture인 SWA는 full attention과 동일해집니다. 이 동등성은 위에서 논의된 관행이 full attention 없는 우리 방법의 진정한 동작을 드러내지 못할 수도 있다는 우려를 제기합니다. 그래서 우리는 오른쪽 패널들에 보이는 것처럼 fine-tuning 후 32K context length의 Books에서도 평가합니다.

Model size에 따른 scaling을 위해, 우리는 단순히 우리의 basic recipe에 있는 다섯 가지 크기에 걸쳐 변화를 줍니다. Training tokens 수에 따른 scaling을 위해, 우리는 model size를 760M로 고정하고, pre-training과 fine-tuning을 위한 training tokens의 수를 변화시킵니다. 구체적으로, pre-training을 위한 우리의 basic number of tokens는 Chinchilla recipe에서 가져왔으며, fine-tuning을 위한 우리의 basic number는 Appendix B에서 논의된 대로 pre-training을 위한 것의 5%입니다. 우리는 5% 비율을 고정하고, pre-training과 fine-tuning을 위해 basic number의 최대 5배까지 실험합니다.

Similar trend to full attention under large budget. 우리는 패널들에 걸쳐 유사한 경향을 관찰합니다:

• Small compute budget 영역에서 training compute가 많아질수록 full attention에 대한 TTT-E2E의 이점은 눈에 띄게 감소합니다.

• 그러나, medium compute budget 영역에서, 파란 선이 비교적 평평하게 유지되는 것으로 나타나듯이 TTT-E2E는 full attention과 유사한 scaling trend를 따릅니다. Model size에 따른 scaling에 대해 여전히 작은 상승(uptick)이 있지만, 우리는 전반적인 경향을 고려할 때 더 큰 models에서는 이 상승이 사라질 것으로 예상합니다.

Model size에 따른 scaling의 경우, regime 변화의 경계는 대략 760M입니다. Training tokens 수에 따른 scaling의 경우, 이 경계는 대략 48B입니다. 우리는 Figure 5에 이 경계들을 점선 수직선으로 표시합니다. Gated DeltaNet이 TTT-E2E와 같은 경향을 따른다는 것을 관찰하는 것은 특히 흥미롭습니다. 우리는 이 관찰에 대해 두 가지 잠재적인 설명을 제공합니다:

• 우리의 방법은 Subsection 2.4에서 설명된 바와 같이 Gated DeltaNet과 유사하게 hybrid RNN으로 해석될 수 있습니다. 우리는 RNNs(고정된 크기의 hidden states를 가진 sequence models)가 training compute에 따른 scaling에 대해 유사한 경향을 공유할 것으로 예상합니다.

• Transformers는 RNNs와 비교하여 불충분한 training compute에서 성능이 저조한 것으로 널리 알려져 있습니다. 우리의 관찰은 large compute에서의 RNNs의 결함보다는 small compute에서의 full attention baseline의 결함으로 해석될 수 있습니다.

전반적으로, 우리의 경험적 관찰들은 TTT-E2E가 large-budget production runs에서 training compute에 따른 scaling에 대해 full attention과 동일한 경향을 생성해야 함을 강력하게 나타냅니다.

Sensitivity to tokenizer and data quality. Scaling 조사 동안, 우리는 최신성이 나타내는 바와 같이 tokenizer와 data quality의 효과에 대한 일화적인 관찰들을 수집했습니다. 구체적으로:

• Llama 2 tokenizer(2023)에서 Llama 3 tokenizer(2024)로 전환하는 것은 3B models에 대해 full attention 대비 우리의 이점을 약 0.01만큼 향상시켰습니다.

• SlimPajama(2023)에서 DCLM(2024)으로 전환하는 것은 우리의 방법이 48B 이후 training tokens 수에 따른 scaling에 대해 full attention과 동일한 경향을 생성할 수 있게 했습니다; FineWebEdu(2024)와의 경향 또한 full attention과 동일합니다. SlimPajama를 사용했을 때, Figure 5의 오른쪽 패널들에서의 우리의 선들은 model size에 따른 scaling을 위한 왼쪽 패널들의 것들과 유사하게 작은 상승을 보였습니다.

이러한 효과들에 대한 포괄적인 조사는 다양한 tokenizers와 datasets에 대해 Figure 5를 재생산하는 것을 수반하며, 이는 우리 논문의 범위를 벗어납니다. 그럼에도 불구하고, 우리의 일화적인 관찰들은 여전히 향후 연구를 위한 시작점을 제공할 수 있습니다. 특히 흥미로운 방향은 self-generated tokens에 대한 TTT인데, 이는 현재 mini-batch의 tokens의 필터링되거나 환어(rephrased)된 버전이거나 이전 mini-batches의 리뷰일 수 있습니다. RNNs의 gating mechanisms가 spurious inputs로부터 hidden states를 보호하고 가치 있는 정보 내의 정보를 더 잘 유지할 수 있다는 것은 널리 알려져 있습니다. 우리는 TTT 동안의 self-generation이 유사한 역할을 할 수 있다고 믿습니다.

3.4 Scaling with Context Length

우리는 첫 페이지의 Figure 1에서 context length에 따른 scaling에 대한 주요 결과들을 제시했습니다. 여기서, 우리는 이러한 실험들의 setup을 논의하고 Figure 6에서 이러한 결과들의 일부에 대한 분석(breakdown)을 제시합니다. 또한, 부록의 Figure 9는 loss $\Delta$s 대신 Figure 1의 loss values를 직접 도표로 나타냅니다.

Figure 1의 실험들을 위해, 우리는 basic recipe에서 가장 큰 model (3B)을 사용합니다. 우리는 또한 pre-training과 fine-tuning 모두에 대해 basic number of tokens의 3배를 사용합니다. 논의된 바와 같이, pre-training을 위한 basic number는 Chinchilla recipe에서 가져왔으며, fine-tuning을 위한 것은 pre-training의 5%입니다. 우리의 이전 실험들과 마찬가지로, 우리는 DCLM에서 하나의 checkpoint를 pre-train한 다음 Books에서 다섯 가지 버전을 fine-tune하는데, 각 context length마다 하나씩 하므로 final results는 Books에서 평가됩니다.

3.4.1 Loss Breakdown by Token Index

Figure 6는 32K와 128K의 두 context lengths에 초점을 맞추고 Figure 1의 해당 결과들을 token index별로 분석합니다; 우리는 Figure 2의 오른쪽 패널을 생성하기 위해 Subsection 2.1에서와 동일한 과정을 따랐습니다. 구체적으로, context length $T$가 주어졌을 때, 각 $t = 1, ..., T$에 대해, 우리는 $x_0, ..., x_{t-1}$을 conditioning하고 $x_t$를 예측하려고 시도하는 next-token prediction task의 test loss를 도표로 나타냅니다. 따라서, context length $T$를 가진 각 method에 대해, Figure 1에서의 test loss는 Figure 6의 해당 곡선에 있는 모든 losses의 평균입니다. 32K에 대한 breakdown은 128K에 대한 것의 subset이 아니라는 점에 유의하는 것이 중요한데, 이는 그것들이 두 개의 서로 다른 models로부터 생성되었기 때문입니다.

우리는 Figure 6의 두 패널 모두에서 다음 관찰들을 합니다:

• TTT-E2E는 전체 context length에 걸쳐 항상 full attention보다 낮은 losses를 달성하는 유일한 method입니다.

• TTT-E2E와 full attention 사이의 test loss 차이는 context window의 끝부분에서 작습니다. Full attention에 대한 TTT-E2E의 집계된 이점은 대부분 앞부분의 tokens에서 나옵니다.

두 관찰이 두 패널 모두에서 동시에 성립한다는 사실은 다소 역설적인 방식으로 특히 흥미롭습니다. 두 번째 관찰의 일부로서, 왼쪽 패널에서 TTT-E2E와 full attention 사이의 차이는 context window의 끝인 $t=32K$ 근처에서 작습니다. 다른 정보가 없다면, 누군가는 곡선들이 128K와 같은 더 큰 context lengths에서 교차할 것이라고 추측할 수도 있습니다. 그러나 이 추측은 오른쪽 패널의 첫 번째 관찰에 의해 단언되듯이 거짓입니다. 128K에 대한 breakdown plot은 추측된 연속이라기보다는 32K에 대한 것의 늘어난(stretched out) 버전을 더 닮았습니다. TTT-E2E가 Figure 1의 context lengths에 걸쳐 full attention에 대해 동일한 이점을 유지한다는 점을 고려하면, 이 stretching effect는 놀랍지 않을 것입니다.

무엇이 TTT-E2E에게 앞부분의 tokens에 대해 full attention보다 이점을 줄까요? 이 이점은 TTT가 첫 번째 (inner-loop) mini-batch에 대해 첫 번째 gradient step을 밟을 때인 $t = 1K$ 이전에도 존재한다는 점에 유의하십시오. 다시 말해, $t = 1K$ 이전에, TTT-E2E와 full attention은 정확히 동일한 computation graph를 가지며 오직 weights만 다릅니다. 그렇다면 왜 TTT-E2E의 weights가 훨씬 더 낮은 losses를 생성할까요?

여기에 직관적인 설명이 있습니다: Full attention의 weights는 context window 내의 모든 미래의 tokens에서 잘하도록 준비해야 합니다. 그러한 작업은 매우 어려울 수 있는데, 모든 가능한 미래에서 잘하는 것은 특정 하나에서 잘하는 model의 capacity를 제한하기 때문입니다. 그러나 TTT-E2E의 weights는 오직 현재의 mini-batch of tokens에서만 잘하면 되는데, TTT가 미래의 tokens를 위해 미래의 weights를 생성할 것이기 때문입니다. 이 더 집중된 작업은 훨씬 더 쉬울 것입니다. 사실, 우리가 Subsection 4.2에서 논의할 TTT의 일반적인 핵심 직관은 현재(present)에 집중하는 것입니다.

3.5 Needle in a Haystack

Section 1에서 논의된 바와 같이 우리 방법의 motivation은 모든 세부 사항을 회상(recall)할 필요 없이 language modeling에서 더 나은 성능을 달성하기 위해 더 긴 context를 사용하는 것이었습니다. 지금까지, 우리는 상세한 회상을 요구하지 않는 evaluations에 초점을 맞췄습니다. 여기서, 우리는 Needle in a Haystack (NIAH)라고 알려진 회상을 위해 명시적으로 설계된 인기 있는 evaluation을 고려합니다: Model은 passage (haystack) 내에서 target string (needle)을 검색해야 하며, 여기서 target string은 나머지 passage와의 명확한 무관성으로 구별됩니다. 구체적으로, 우리는 128K context length에서 fine-tune된 모든 3B models를 RULER의 세 가지 NIAH tasks에서 평가합니다.

Table 2에서, 우리는 full attention을 갖춘 Transformer가 우리 것을 포함한 다른 methods보다, 특히 long context에서 극적으로 성능이 우수함을 관찰합니다. 우리의 이전 subsections의 발견들과 결합된 이 관찰은 full attention의 강점이 거의 손실 없는 회상(lossless recall)에 있다는 직관을 뒷받침합니다. 이 강점은 cache에 있는 모든 이전 tokens의 keys와 values에 attend하는 self-attention의 설계에 내재되어 있습니다. 대조적으로, 우리 방법의 핵심 메커니즘은 target string과 같이 겉보기에 무관한 세부 사항들을 생략하는 압축(compression)입니다.

3.6 Decoding Long Sequences

지금까지, 우리의 모든 evaluations는 model이 십여 개의 tokens보다 많이 decode하지 않을 것을 요구했습니다. Subsection 2.3의 끝에서 논의된 바와 같이, decoded tokens가 TTT mini-batch를 채우면, TTT-E2E는 이 decoded tokens 배치에 대해 gradient step을 밟습니다. Test time에서의 이러한 “self-training” 방법이 long sequences를 decoding하는 데 효과가 있을까요?

실제로, long sequences decoding을 요구하는 시나리오들은 일반적으로 instruction fine-tuning 이후나 reinforcement learning 동안, 예를 들어 model이 긴 chains of thought를 생성할 때 발생합니다. 따라서, 위의 두 stages 없이 base models를 현실적인 방식으로 평가하는 것은 본질적으로 어렵습니다. 이 두 stages는 우리 논문의 범위를 벗어나므로, 우리는 Subsection 3.4에서 훈련한 3B base models를 평가하기 위해 최선을 다합니다.

Figure 7의 evaluation을 위해, 우리는 evaluator로 Qwen-3-8B-Base를 사용합니다. 우리의 models는 Books에서 훈련되었으므로, 우리는 그것들의 context windows를 Books의 8K tokens로 prefill하고, 또 다른 8K tokens를 continuation으로 decode한 다음, 연결된 16K sequence에 대한 Qwen-8B의 loss (log likelihood)를 token index별로 도표로 나타냅니다. Figure 6가 x축에 log scale을 사용하는 반면, 여기 Figure 7은 linear scale을 사용하여 prefill과 decode에 대한 경향을 쉽게 비교할 수 있게 합니다. 이 evaluation의 추가 세부 사항은 Appendix D에 제공됩니다.

우리의 이전 관찰들과 유사하게, TTT-E2E는 이 제한된 evaluation에서 full attention보다 낮은 Qwen loss를 달성합니다. 게다가, 우리는 생성된 텍스트의 약 20개 샘플을 주의 깊게 검사했고 합리적임을 발견했습니다. 두 methods 모두에 대해, Qwen loss는 prefill과 decode 사이의 경계에서 급격히 증가하다가 다시 점진적으로 감소합니다. 이러한 동작은 아마도 Qwen이 처음에는 평가되는 method의 generation style에 익숙하지 않지만, context window 내에 더 많은 생성된 콘텐츠가 축적됨에 따라 점차 적응하기 때문에 발생할 것입니다.

3.7 Computational Efficiency

Figure 1에서, 우리는 baselines의 것과 비교하여 우리의 inference latency, 구체적으로 prefill latency를 제시했습니다. 여기서, 우리는 prefill latency를 측정하기 위한 setup을 논의하고, computational efficiency가 중요한 두 가지 추가 축인 decode와 training을 고려합니다. 특히, 우리는 training latency를 현재 구현의 중요한 한계점으로 강조하고 이를 개선하기 위한 두 가지 잠재적인 방향을 논의합니다.

Setup for prefill latency. Figure 1의 오른쪽 패널에 있는 각 method에 대해, 우리는 왼쪽 패널의 해당 3B model을 가져와 하나의 H100에서 prefill latency를 측정했습니다. 우리는 또한 Appendix C에서 논의된 바와 같이 PyTorch baselines의 inference latency를 최적화하기 위해 추가적인 조치를 취했습니다. Gated DeltaNet을 따라, latency experiments는 (outer-loop) batch 당 일정한 수의 tokens (128K)로 수행됩니다. 예를 들어, 128K context length에서 각 batch는 하나의 sequence를 포함하고, 8K에서는 각 batch가 16개의 sequences를 포함합니다.

TTT-E2E only uses standard infrastructure. Test time에, TTT-E2E는 단순히 regular Transformer를 훈련하기 위해 최적화된 standard infrastructure를 사용할 수 있습니다. 구체적으로, 우리의 hidden state는 regular MLP layers의 형태를 취하므로, custom kernel 없이 standard tools를 사용하여 GPU들에 걸쳐 shard될 수 있습니다. 대조적으로, 이전 연구는 그들의 hidden states를 GPU 내부의 개별 칩들에 맞게 넣어야 하므로 hidden state size가 상당히 제한됩니다. 예를 들어, TTT-KVB는 LoRA로 state size를 줄여야 하는 반면, Mamba 2 및 Gated DeltaNet과 같은 다른 이전 연구는 linear hidden state를 사용하고 효율적인 memory I/O를 위해 custom kernels를 작성해야 합니다.

Decode latency. Subsection 2.3의 끝에서 논의된 바와 같이, 우리의 방법은 decoded tokens가 TTT mini-batch를 완전히 채울 때까지 TTT를 수행하지 않습니다. 따라서 full batch에 도달하기 전까지, 우리의 decode latency는 SWA를 갖춘 regular Transformer의 것과 같습니다. 일단 full batch를 갖게 되면, 우리는 다음 tokens 배치를 decode하기 전에 TTT step이 필요하며, 이 TTT step에 대한 우리의 latency는 prefill에 대한 것과 같습니다. 전체적으로, multiple batches의 long sequence를 decoding하기 위한 우리의 latency는 단순히 위의 두 latencies의 합입니다: SWA decode의 것과 우리의 prefill의 것. 둘 다 쉽게 이용 가능하므로, 우리는 TTT-E2E의 decode latency에 대한 별도의 측정값을 보고하지 않습니다.

Setup for training latency. 우리의 훈련 대부분은 GB200s에서 수행되었습니다. 우리의 baselines 중 다수가 GB200s (Blackwell)를 위해 작성된 custom kernels를 가지고 있지 않으므로, 우리는 baselines에 대한 공정성을 위해 H200 (Hopper)에서 training latency를 벤치마크합니다. Prefill에 대한 프로토콜을 따라, 우리는 context length에 관계없이 batch 당 일정한 수의 tokens (128K)를 사용합니다.

Training latency is a limitation. Training time에, TTT-E2E는 gradients of gradients를 취하는데, 이는 regular Transformer 훈련과 비교하여 훨씬 덜 최적화된 절차입니다. Figure 8의 왼쪽 패널에서 보여지듯이, 우리의 training latency는 128K context length에서 full attention보다 1.2배 빠르지만, 8K에서는 3.4배 느립니다. Training compute의 대부분이 일반적으로 short context에서의 pre-training에 소비되므로, 우리 현재 구현의 training latency는 여전히 중요한 한계점으로 남습니다. 비록 오른쪽 패널에서 보여지듯이 token 당 FLOPs 수는 일정하게 유지되지만, 우리의 latency는 8K와 32K 사이에서 증가한다는 점에 유의하십시오. 이 경향은 우리가 context length가 $T$일 때 $\log(T)$의 인수로 시간 경과에 따른 gradient checkpointing의 양을 증가시켜야 하기 때문에 발생합니다.

Directions for faster training. 전체적인 training time을 개선하기 위한 두 가지 방향이 있습니다:

• 우리의 현재 구현은 gradients of gradients를 지원하지 않기 때문에 training time에 cuDNN FlashAttention을 사용할 수 없습니다. Custom attention kernel은 우리의 hardware utilization을 유의미하게 개선하고, 잠재적으로 시간 경과에 따른 gradient checkpointing으로 인한 바람직하지 않은 경향을 제거할 것입니다.

• 우리는 TTT-E2E의 훈련이 TTT 없는 pre-trained Transformer로부터 초기화될 수 있다고 믿습니다 – 이는 RNNs에 대한 이전 연구에서 자주 채택되는 기술입니다. 이 실용적인 기술은 TTT-E2E가 전체 training compute의 작은 부분만을 차지하게 하여, training latency의 부정적인 효과를 최소화합니다.

우리는 이러한 방향들을 향후 연구로 남겨둡니다.

 

 

 

Section 3.1 핵심 정리 노트 (Target: AI Researchers)

이 섹션은 TTT-E2E의 성능을 검증하기 위한 Experimental Design과 Fair Comparison을 위한 Baseline 설정을 다룹니다.

  • Experimental Goal & Scope
    • Large-scale production training을 모사할 수 있는 Small & Medium scale (125M ~ 3B parameters) 환경에서의 검증.
    • Training Strategy (2-Stage):
      1. Pre-training: Short context ($8K$)로 general knowledge 학습.
      2. Long-context Extension (Fine-tuning): $128K$까지 context length를 확장.
      • (Note: 실제 Large-scale에서는 더 많은 stage가 있지만, 연구 목적상 2단계로 단순화함)
  • Datasets Selection
    • Pre-training: DCLM-Baseline (Common Crawl subset).
      • Preprocessing: $8K$보다 짧은 문서는 버리고(discard), 남은 것 중 random sampling.
    • Fine-tuning (Extension): Books dataset.
      • Reasoning: DCLM의 $128K$ 이상 데이터는 품질이 낮아(low quality), Long-context standard인 Books를 채택.
  • Baselines (State-of-the-Art Comparisons)
    • 모든 Sliding Window 기반 모델은 Window size $k=8K$ 로 통일.
    1. Transformer with full attention: Gold standard (Upper bound).
    2. Transformer with SWA (Sliding-Window Attention):
      • TTT-E2E의 Base Architecture.
      • 중요: Pre-training context가 $8K$이므로, Fine-tuning 전까지는 Full attention과 SWA는 동일한 모델임 ($k=8K$일 때 Full attention = SWA).
    3. Hybrid SWA and full attention (5:1): Gemma 스타일 (5 layers SWA + 1 layer Full).
    4. Mamba 2: Hybrid of Mamba 2 & SWA layers.
    5. Gated DeltaNet: Hybrid of Gated DeltaNet & SWA layers (Linear Attention 계열).
    6. TTT-KVB: Prior TTT work. Layer-wise reconstruction loss를 사용하는 방식 (Section 2.4의 출발점).
  • Implementation Details
    • Configuration: GPT-3 및 Mamba의 standard config를 따름 (Appendix B 참조).
    • Optimization: Transformer baseline에 대해 Grid search를 수행하여 최적의 Hyper-parameter 도출.
    • Framework: Baselines 1~3 및 TTT-E2E는 JAX로 구현, 나머지는 저자들의 Official Code 사용.

쉬운 설명 : "공정한 싸움을 위한 링(Ring) 만들기"

이 섹션은 TTT-E2E가 진짜로 성능이 좋은지 확인하기 위해 **"어떤 조건에서, 누구와 싸웠는지"**를 설명하는 부분입니다.

  1. 체급 설정 (Model Size): 거대 언어 모델(GPT-4급)을 처음부터 학습시키는 건 너무 비싸고 오래 걸리니, 중소형 모델(125M~3B) 사이즈로 실험했습니다. 여기서 이기면 거대 모델에서도 이길 것이라고 가정하는 것이죠.
  2. 커리큘럼 (Training Stages):
    • 초등학교 (Pre-training): 먼저 8K 길이의 짧은 글(DCLM)로 일반 상식을 공부시킵니다.
    • 대학교 (Fine-tuning): 그 다음 128K 길이의 아주 긴 책(Books)을 읽히면서 "긴 글을 소화하는 법"을 가르칩니다.
  3. 경쟁자들 (Baselines):
    • Full Attention: "기억력 끝판왕". 모든 걸 다 기억하지만 머리가 터질 듯이 느린 친구. (우리의 목표는 얘만큼 똑똑하면서 속도는 빠른 것)
    • SWA: "건망증 친구". 최근 8K만 기억하고 나머진 까먹는 친구. (TTT-E2E는 이 친구의 뇌 구조를 그대로 쓰면서 학습법만 바꾼 것)
    • Mamba 2, DeltaNet: "요즘 뜨는 신예들". RNN 계열의 최신 모델들.
    • TTT-KVB: "옛날 방식의 TTT".

저자들은 특히 **"우리는 SWA랑 뇌 구조(Architecture)가 똑같은데, 학습 방법(Algorithm)만 바꿔서 천재(Full Attention)를 이겼다"**는 것을 보여주기 위해 이 셋업을 짰습니다.

 

 

 

더보기

Figure 4: Ablations on Hyper-Parameters (하이퍼파라미터 검증)

이 그림은 TTT-E2E의 성능에 영향을 미치는 세 가지 주요 요소(Hyper-parameter)를 실험한 결과입니다. (Section 3.2 참조)

  • Left (Sliding Window Size $k$):
    • 내용: Window Size $k$가 커질수록(1K → 8K) 성능(Loss)이 어떻게 변하는지 보여줍니다.
    • 결과: SWA(녹색), Gated DeltaNet(갈색), TTT-E2E(파란색) 모두 $k$가 클수록 성능이 좋아집니다.
    • 결론: 성능과 효율성을 고려하여 기본값을 8K로 설정했습니다.
  • Middle (TTT mini-batch size $b$):
    • 내용: TTT를 수행할 때 한 번에 업데이트하는 토큰 묶음의 크기($b$)에 따른 성능 변화입니다.
    • 결과: $b$가 커질수록 성능이 나빠집니다(Loss 증가). 이는 배치 내부에서 업데이트가 즉각적으로 반영되지 않는 문제(Intra-batch amnesia) 때문입니다. 반면 너무 작으면 하드웨어 효율이 떨어집니다.
    • 결론: 효율과 성능의 타협점으로 1K를 선택했습니다.
  • Right (Number of Layers Updated):
    • 내용: 전체 레이어 중 몇 개의 레이어를 TTT로 업데이트할 것인가에 따른 확장성(Scaling) 실험입니다.
    • 결과: 1개(회색)나 3개(빨간색) 레이어만 업데이트하면 긴 문맥(Context)에서 성능 유지가 안 됩니다(Loss $\Delta$ 상승). 반면, 6개(파란색, 전체의 1/4) 혹은 12개(보라색)를 업데이트하면 Full Attention(주황색)과 유사한 패턴으로 성능이 유지됩니다.
    • 결론: 효율성을 위해 마지막 1/4 레이어만 업데이트하기로 결정했습니다.

 

Section 3.2 ~ 3.2.1 핵심 정리 노트 (Target: AI Researchers)

이 섹션은 TTT-E2E의 성능을 결정짓는 핵심 Hyper-parameter들에 대한 Ablation Study를 통해, 모델의 Empirical Intuition을 확립하는 과정입니다.

  • Sliding Window Size $k$ (vs. Full Attention)
    • Observation: $k$가 커질수록 SWA, Gated DeltaNet, TTT-E2E 모두 성능이 향상됨.
    • Orthogonal Improvement (Key Insight):
      • Pre-training context와 동일한 $k=8K$ 설정 시, SWA는 사실상 Full Attention과 동일함.
      • 이 상태에서도 TTT-E2E는 Full Attention 대비 Test Loss를 0.018 개선함.
      • $k$가 증가해도 SWA와 TTT-E2E 간의 성능 격차는 일정하게 유지됨.
      • Conclusion: TTT는 단순히 SWA의 짧은 Context를 보완하는 역할(Compensating)이 아니라, Attention 메커니즘과는 별개로 작동하는 직교하는(Orthogonal) 성능 향상 요인임.
  • TTT Mini-batch Size $b$
    • Trade-off:
      • Small $b$: Frequent updates로 성능 좋음 (Intra-batch amnesia 최소화).
      • Large $b$: Parallelism 및 Stability 유리.
    • Result: $b$가 커질수록 성능이 Significantly 하락함. 반면 $1K$ 미만은 하드웨어 효율성이 급감함.
    • Decision: $b=1K$를 Default로 선정.
  • Architecture vs. Algorithm (Role of Modifications)
    • Question: Figure 3에서 적용한 아키텍처 수정(Sliding window, Dual MLP 등)이 성능 향상의 주원인인가?
    • Verification: TTT 과정을 끄고($b$를 매우 크게 설정하여 업데이트 안 함) 아키텍처만 유지했을 때의 성능 측정.
    • Result: TTT 없는 TTT-E2E(2.825) $\approx$ Vanilla Full Attention(2.827).
    • Conclusion: 아키텍처 수정은 보조적인 역할일 뿐, 성능 향상의 핵심은 Test-Time Training 알고리즘 자체에 있음.
  • Number of Layers Updated (Compression Capacity)
    • Hypothesis: 업데이트되는 Layer의 수는 Context 정보를 압축하여 저장할 수 있는 Storage Capacity를 결정함.
    • Experiment: Last 1 layer, 3 layers(1/8), 6 layers(1/4), 12 layers(1/2) 업데이트 비교.
    • Critical Finding:
      • 1 or 3 layers: Context length가 길어질수록 Scaling Law를 따르지 못하고 성능이 무너짐 (Full Attention의 Scaling trend를 따라가지 못함).
      • 6 layers (1/4): Full Attention과 유사한 Scaling trend를 보임.
      • 12 layers (1/2): 1/4 대비 유의미한 성능 차이 없음.
    • Decision: 효율성을 위해 Last 1/4 layers만 업데이트하는 것을 표준으로 채택.

쉬운 설명 : "최적의 뇌 개조 전략 찾기"

이 섹션은 "시험 시간에 뇌(Weights)를 고쳐서 공부하는 방법(TTT)"을 쓸 때, 구체적으로 어떻게 고쳐야 가장 효율적인지를 실험한 부분입니다.

1. 안경(Window Size)이 이미 좋은데 뇌 개조가 필요한가?

  • 실험: 시력이 아주 좋아서 모든 지문을 다 볼 수 있는 상태(Full Attention)에서도 TTT를 하는 게 의미가 있을까요?
  • 결과: 네, 의미가 있습니다. 시력이 좋은 것과 별개로, 뇌를 학습시키면 점수가 더 오릅니다. 즉, TTT는 단순히 눈이 나쁜(SWA) 모델을 돕는 보청기가 아니라, 이미 똑똑한 모델을 더 천재로 만드는 부스터입니다.

2. 벼락치기의 호흡 (Batch Size)

  • 실험: 책을 읽을 때 몇 페이지마다 멈춰서 뇌를 고치는 게 좋을까요?
  • 결과: 책을 너무 많이 읽고 나서(Large Batch) 한 번에 고치려고 하면, 읽는 도중에 앞 내용을 까먹습니다. 그렇다고 한 줄 읽고 고치면(Small Batch) 너무 산만합니다. 1,000 토큰(약 1K) 정도 읽고 고치는 게 딱 좋습니다.

3. 뇌 구조가 중요한가, 공부법이 중요한가?

  • 실험: TTT를 안 하고 모델 구조만 이 논문처럼 바꾼다면 어떨까요?
  • 결과: 아무 효과 없습니다. 일반 모델과 똑같습니다. 즉, 특별한 구조 덕분이 아니라, 실시간으로 공부하는 방식(TTT) 덕분에 성능이 좋은 것입니다.

4. 뇌의 어느 부분을 고칠까? (Layers Updated)

  • 실험: 뇌 전체를 다 고칠까요, 아니면 일부만 고칠까요?
  • 결과: 뇌의 아주 일부(1개 레이어)만 고치면 용량이 부족해서 긴 책 내용을 다 못 담습니다. 하지만 뒤쪽 1/4 정도만 고쳐도 충분히 모든 내용을 담을 수 있습니다. 그 이상 고치는 건 낭비입니다.

 

 

 

 

 

 

 

더보기

Figure 5: Scaling with Training Compute (학습 연산량에 따른 확장성)

모델의 크기(Model Size)와 학습 데이터 양(Training Tokens)을 늘렸을 때, TTT-E2E가 Full Attention만큼 성능이 향상되는지 확인한 그래프입니다. (Section 3.3 참조)

  • Left Panels (Scaling with Model Size):
    • 모델 크기를 125M에서 3B로 키웠을 때의 결과입니다.
    • TTT-E2E(파란색)는 Full Attention(주황색, 기준선 0)보다 항상 더 좋은 성능($y < 0$)을 보입니다.
    • 큰 모델(3B)로 갈수록 격차가 약간 줄어들지만, Full Attention과 유사한 추세(평행한 기울기)를 유지합니다.
  • Right Panels (Scaling with Training Tokens):
    • 학습 토큰 수를 16B에서 80B로 늘렸을 때의 결과입니다.
    • 점선 오른쪽(Large budget) 영역을 보면, TTT-E2E(파란색)가 Full Attention(주황색)과 평행하게 가거나 더 좋은 성능을 유지함을 알 수 있습니다.
    • 결론: TTT-E2E는 대규모 학습 환경에서도 Full Attention과 동일한 **Scaling Law(확장 법칙)**를 따릅니다.

 

더보기

Figure 6: Loss Breakdown by Token Index (토큰 위치별 성능 분석)

긴 문맥(32K, 128K) 내에서 각 토큰 위치($t$)별로 모델이 얼마나 잘 예측하는지를 뜯어본 그래프입니다. (Section 3.4.1 참조)

  • Left (32K Context) & Right (128K Context):
    • X축: 토큰의 위치 (Log scale). 오른쪽으로 갈수록 문맥의 뒷부분입니다.
    • Y축: Loss (낮을수록 좋음).
  • 핵심 결과:
    • TTT-E2E(파란색) 선이 가장 아래에 위치합니다. 즉, 문맥의 처음부터 끝까지 Full Attention(주황색)보다 항상 더 잘 예측합니다.
    • 특히 **초반부(Earlier tokens)**에서 Full Attention보다 성능 이득이 큽니다. 이는 TTT가 "미래의 모든 토큰을 대비해야 하는" Full Attention과 달리, "현재 배치만 잘하면 되는" 더 쉬운 목표를 가지기 때문으로 해석됩니다.
    • 128K 그래프는 32K 그래프가 단순히 늘어난 형태와 유사하며, 긴 문맥에서도 성능이 무너지지 않음을 보여줍니다.

 

 

Section 3.3 핵심 정리 노트 (Target: AI Researchers)

이 섹션은 TTT-E2E가 Scaling Law를 따르는지, 즉 Model Size와 Training Compute가 증가함에 따라 Full Attention과 동일한 기울기(Trend)로 성능이 향상되는지를 검증합니다.

  • Evaluation Protocol Strategy
    • Standard Practice (Pre-training eval): Pre-training 직후 DCLM(8K context)에서 평가.
      • Issue: Pre-training 단계($8K$)에서는 TTT-E2E의 Base Architecture(SWA)와 Full Attention이 동일함. TTT의 진가를 확인하기 어려움.
    • Adjusted Practice (Fine-tuning eval): Fine-tuning 후 **Books (32K context)**에서 평가.
      • Full Attention이 아닌 TTT-E2E 고유의 동작을 확인하기 위한 필수적인 설정.
  • Scaling Results (Compute Budget & Model Size)
    • Small Compute Budget Regime:
      • 초기에는 TTT-E2E가 Full Attention 대비 큰 이득(Loss $\Delta$)을 보이지만, 연산량이 늘어날수록 이득이 감소함.
    • Medium/Large Compute Budget Regime (Critical):
      • Model Size: 760M 이상.
      • Training Tokens: 48B 이상.
      • 이 지점부터는 TTT-E2E가 Full Attention과 **유사한 Scaling Trend(평행선)**를 유지함.
      • Gated DeltaNet (RNN Baseline) 또한 동일한 Trend를 보임.
  • Interpretation of Trends
    1. Hybrid RNN Hypothesis: TTT-E2E는 본질적으로 Hybrid RNN(Section 2.4 유도 과정 참조)으로 볼 수 있으며, RNN 계열이 공유하는 Scaling 특성일 수 있음.
    2. Transformer's Data Hunger: Small compute에서 Full Attention의 성능이 낮은 것은, Transformer가 RNN보다 Inductive Bias가 적어 학습에 더 많은 데이터/연산이 필요하기 때문일 수 있음 (즉, TTT가 나빠지는 게 아니라 Full Attention이 정상 궤도에 오르는 과정).
  • Sensitivity Analysis (Anecdotal Evidence)
    • Tokenizer Impact: Llama 2 $\rightarrow$ Llama 3 Tokenizer 교체 시 3B 모델에서 성능 이득 확대.
    • Data Quality Impact: SlimPajama $\rightarrow$ DCLM/FineWebEdu 교체 시 Scaling Trend가 Full Attention과 더욱 완벽하게 일치.
      • Low quality data에서는 약간의 Uptick(성능 정체) 현상이 관찰되었으나 High quality data에서 해소됨.
    • Future Direction: RNN의 Gating mechanism이 노이즈를 필터링하듯, TTT 과정에서 Self-generated tokens를 사용하여 데이터 품질을 스스로 제어하는 연구 가능성 제시.

쉬운 설명 : "천재는 대기만성형이다?"

이 섹션은 **"이 모델(TTT), 덩치를 키우고 공부를 많이 시켜도 계속 똑똑해질까?"**라는 의문을 해결합니다.

  1. 초반 (공부량이 적을 때):
    • TTT 모델은 조금만 공부해도 성적이 쑥쑥 잘 나옵니다. 반면, 기존의 천재 모델(Full Attention Transformer)은 발동이 늦게 걸려서 초반에는 TTT보다 못합니다.
  2. 중반 이후 (공부량이 많을 때):
    • 기존 천재 모델이 드디어 감을 잡고 무서운 속도로 똑똑해지기 시작합니다. 이때 많은 사람들이 "아, 역시 덩치가 커지면 결국 Transformer가 짱이고, TTT 같은 새로운 시도들은 한계(Wall)에 부딪히지 않을까?"라고 걱정합니다.
    • 하지만 실험 결과, TTT 모델도 천재 모델과 똑같은 속도로 계속 똑똑해집니다. (그래프가 평행하게 감)
    • 즉, "공부를 많이 시키면 효율이 떨어지지 않을까?" 하는 걱정은 안 해도 됩니다.
  3. 데이터 품질의 중요성:
    • 좋은 교재(DCLM, FineWebEdu)로 공부시키면 TTT가 Transformer와 완벽하게 똑같은 성장 속도를 보입니다.
    • 오히려 교재가 나쁠 때 TTT가 잠깐 주춤하는 경향이 있는데, 저자들은 TTT가 스스로 "이건 나쁜 정보야"라고 거르는 능력(Gating)을 키우면 해결될 것이라고 봅니다.

 

 

 

 

 

 

 

 

 

더보기

Table 2: S-NIAH Performance (Needle In A Haystack 성능)

"건초 더미에서 바늘 찾기(Pass-key retrieval)"와 같이, 긴 문맥 속에서 특정 단어를 정확히 기억해내는 능력(Recall)을 평가한 표입니다. (Section 3.5 참조)

  • 결과:
    • Full Attention: 거의 모든 길이(8K~128K)에서 **1.00(완벽함)**에 가까운 점수를 기록합니다.
    • TTT-E2E 및 기타 RNN류: 문맥이 길어질수록 점수가 0.06, 0.05 등으로 처참하게 떨어집니다.
  • 해석:
    • 이것은 논문의 주장과 일치합니다. Full Attention은 모든 정보를 다 저장(Lossless)하므로 검색에 유리하지만, TTT는 정보를 압축(Compression)하므로 "지엽적인 세부 사항(Needle)"은 잊어버립니다.
    • 하지만 앞선 실험들(Figure 1, 6)에서 보듯, 일반적인 언어 모델링(Language Modeling) 성능은 TTT가 더 좋습니다. 즉, "단어 찾기"는 못해도 "글 쓰기/이해하기"는 더 잘한다는 뜻입니다.

Section 3.4 핵심 정리 노트 (Target: AI Researchers)

이 섹션은 TTT-E2E가 긴 문맥(Context)에서 실제로 Full Attention 대비 어떤 우위를 가지는지, 그리고 그 성능 향상의 원인이 무엇인지를 Token-level Breakdown을 통해 분석합니다.

  • Setup
    • Model: 3B Parameters.
    • Training: 3x Chinchilla tokens budget.
    • Eval: 32K 및 128K Context Length에서 Fine-tuning 후 평가.
  • Key Observation: Consistent Superiority
    • All-range Dominance: TTT-E2E는 Context의 길이($T$)와 무관하게, $t=1$부터 $t=T$까지 모든 구간에서 Full Attention보다 낮은 Loss를 기록함 (Figure 6).
    • Source of Advantage: Full Attention 대비 성능 이득(Aggregated Advantage)의 대부분은 **초반부 토큰(Earlier Tokens)**에서 발생함.
    • Gap Narrowing but Preserved: Context의 끝부분($t \approx 32K$)으로 갈수록 Full Attention과의 Loss 차이가 줄어들지만, 128K 실험에서도 역전되지 않고 이 경향성이 그대로 늘어남(Stretched).
  • Theoretical Insight: "Focus on the Present"
    • Question: $t < 1K$ 구간(첫 번째 Gradient Step 이전)에서는 TTT-E2E와 Full Attention의 Computation Graph가 동일한데, 왜 TTT-E2E의 초기 Loss가 더 낮은가?
    • Hypothesis:
      • Full Attention's Burden: Weights가 미래의 모든 시점($t=1$부터 $t=T$까지)에 대해 일반화(Generalize)되어야 함. 모든 가능성을 대비해야 하므로 특정 시점에 최적화되기 어려움.
      • TTT-E2E's Freedom: 현재의 Weights는 현재의 Mini-batch만 잘 처리하면 됨. 미래의 토큰은 미래의(업데이트된) Weights가 처리할 것이기 때문.
    • Conclusion: TTT 메커니즘은 모델이 "먼 미래를 미리 걱정할 필요 없이" **현재에 집중(Specialization)**할 수 있게 해주며, 이것이 초기 구간부터 성능 우위를 가져오는 핵심 원인임.

쉬운 설명 : "오늘만 사는 모델 vs 평생을 걱정하는 모델"

1. 초반부터 압도하는 이유

그래프를 뜯어보니, TTT 모델은 문맥의 뒷부분보다 **앞부분(초반)**에서 점수를 엄청나게 따고 들어갑니다. 심지어 아직 학습(Update)을 시작도 안 한 맨 처음 구간에서도 일반 모델보다 점수가 좋습니다. 왜 그럴까요?

2. 짐 쌀 때의 마음가짐 차이

  • 일반 모델 (Full Attention): "나는 한 번 배낭을 싸면($W$), 여행이 끝날 때까지 절대 못 바꿔. 그러니까 비가 올 때, 눈이 올 때, 더울 때를 다 대비해서 짐을 완벽하게 싸야 해."
    • $\rightarrow$ 모든 상황을 대비하려다 보니, 정작 당장 필요한 물건을 깊이 있게 챙기지 못합니다. (부담감 $\uparrow$, 효율 $\downarrow$)
  • TTT 모델: "어차피 여행 중간중간에 배낭을 다시 쌀 수 있잖아($W_t$). 지금은 당장 내일 필요한 것만 챙기자. 겨울옷은 겨울이 오면 그때 챙기면 돼."
    • $\rightarrow$ 미래의 걱정은 미래의 나에게 맡기고, 지금 당장(Present) 필요한 것에만 100% 집중합니다.

3. 결론

TTT 모델은 "나중에 업데이트할 수 있다"는 믿음 덕분에, 초기 가중치(Initial Weights)가 먼 미래까지 커버해야 한다는 부담을 내려놓고 현재 닥친 문제에만 최적화될 수 있습니다. 이것이 바로 TTT가 처음부터 끝까지 더 똑똑한 이유입니다.

 

 

 

 

 

 

 

더보기

Figure 7: Decoding Long Sequences (긴 문장 생성 능력)

모델이 직접 긴 글(8K 토큰)을 생성(Decoding)하게 한 뒤, 그 글의 품질을 평가(Qwen-8B 모델 사용)한 그래프입니다. (Section 3.6 참조)

  • X축: 0~8K는 주어진 문맥(Prefill), 8K~16K는 모델이 생성한 부분(Decode).
  • Y축: 생성된 글의 Loss (낮을수록 품질이 좋음).
  • 결과:
    • 8K 지점(점선) 이후 생성 단계에서 **TTT-E2E(파란색)**가 Full Attention(주황색)보다 더 낮은 Loss를 유지합니다.
    • 즉, TTT 방식의 "Self-training (생성하면서 학습하기)"이 긴 글을 생성할 때도 효과적임을 보여줍니다.

 

 

Section 3.5 & 3.6 핵심 정리 노트 (Target: AI Researchers)

이 섹션은 TTT-E2E의 **명확한 한계점(Recall)**과 **새로운 가능성(Generation)**을 대조적으로 보여줍니다.

3.5 Needle in a Haystack (The Limitation: Recall vs. Compression)

  • Evaluation Goal: "압축(Compression)" 기반 모델이 "회상(Recall)" 중심의 과제를 얼마나 잘 수행하는지 검증.
    • Task: RULER 벤치마크의 NIAH(Needle In A Haystack). 긴 문맥(Haystack) 속에 숨겨진 무관한 문자열(Needle, 예: UUID)을 찾아내는 과제.
  • Critical Observation:
    • Full Attention: Context length($128K$)에 상관없이 거의 완벽한 성능($\approx 1.00$) 유지.
      • Reason: KV Cache에 모든 Key-Value를 명시적으로 저장하므로 Lossless Recall이 가능함.
    • TTT-E2E & RNNs: Context가 길어질수록 성능이 급격히 하락하여 $128K$에서는 거의 실패($< 0.10$).
      • Reason: TTT의 핵심 메커니즘은 Compression임. UUID와 같이 문맥상 중요해 보이지 않는(Irrelevant) 세부 정보는 가중치 업데이트 과정에서 손실됨.
  • Conclusion: 단순 검색/회상(Retrieval) 능력에 있어서는 Full Attention이 압도적이며, 이는 압축 기반 모델(TTT)의 구조적 한계임.

3.6 Decoding Long Sequences (The Strength: Self-Training at Test Time)

  • Evaluation Goal: TTT가 프롬프트 처리(Prefill)뿐만 아니라, 긴 문장 생성(Decoding) 시에도 유효한지 검증.
  • Mechanism: "Self-Training"
    • 모델이 생성한 토큰들이 TTT mini-batch($1K$)를 채우면, 자신이 생성한 텍스트에 대해 Gradient Step을 수행하여 Weights를 업데이트함.
  • Experiment Setup:
    • Model: 3B trained on Books.
    • Task: $8K$ Prefill (Books) + $8K$ Decode (Continuation).
    • Metric: Qwen-3-8B-Base를 Evaluator로 사용하여 생성된 텍스트의 Log-likelihood 측정.
  • Results:
    • Superior Quality: 생성 단계(Decode)에서 TTT-E2E가 Full Attention보다 더 낮은 Qwen Loss(더 나은 품질)를 기록함.
    • Behavior: 두 모델 모두 Prefill $\rightarrow$ Decode 전환 시점에서 Loss가 급증했다가 감소하는 경향을 보임(Evaluator가 생성 스타일 적응).
  • Conclusion: TTT는 긴 글을 생성할 때, 자신이 쓴 글을 통해 실시간으로 학습하며 성능을 유지/향상시킬 수 있음.

쉬운 설명 : "건망증 천재 작가"

1. 전화번호부 외우기 (Section 3.5 - 한계점)

  • 상황: 책 한 권을 다 읽고 나서, "34페이지 7번째 줄에 있던 이상한 숫자 코드가 뭐였지?"라고 물어봅니다.
  • 일반 모델 (Full Attention): 책을 통째로 복사해서 가방(Cache)에 넣어뒀기 때문에, 가방을 뒤져서 100% 정확하게 찾아냅니다.
  • TTT 모델: 책의 "내용"과 "흐름"을 이해해서 뇌(Weights)에 저장했습니다. 그래서 "주인공이 왜 울었는지"는 알지만, "의미 없는 숫자 코드" 같은 건 중요하지 않다고 생각해서 까먹어버렸습니다.
    • $\rightarrow$ 결론: 팩트 체크나 단순 검색은 일반 모델이 훨씬 잘합니다.

2. 소설 이어 쓰기 (Section 3.6 - 강점)

  • 상황: 앞부분 줄거리를 주고 "뒷내용을 100페이지 더 써봐"라고 시킵니다.
  • 일반 모델: 글을 쓰면서 앞 내용을 참고만 합니다. 글을 쓸수록 피로도가 쌓입니다.
  • TTT 모델: 자기가 글을 쓰면서 **"내가 방금 쓴 내용도 공부(Update)"**합니다. 10페이지를 쓰고 나면, 그 10페이지를 쓴 경험치로 뇌를 업그레이드해서 11페이지를 더 잘 씁니다.
    • $\rightarrow$ 결론: 긴 호흡의 창작이나 글쓰기는 TTT 모델이 더 자연스럽고 퀄리티 있게 잘합니다.

 

 

 

 

 

 

 

더보기

Figure 8: Training Efficiency (학습 효율성)

학습(Training) 단계에서의 속도와 연산량을 비교한 그래프입니다. (Section 3.7 참조)

  • Left (Latency): 학습 속도 (초/1k 토큰).
    • Full Attention(주황색)은 길이가 길어지면 속도가 기하급수적으로 느려집니다(치솟는 그래프).
    • TTT-E2E(파란색)는 **상수 시간(Constant)**에 가깝지만, 초기 높이가 높습니다. 즉, 짧은 문맥(8K~32K)에서는 Full Attention보다 느리지만, **매우 긴 문맥(128K 이상)**에서는 더 빠릅니다. 이는 이중 미분(Gradients of gradients) 계산 오버헤드 때문입니다.
  • Right (FLOPs): 연산량.
    • Full Attention은 $O(T^2)$로 연산량이 폭증합니다.
    • TTT-E2E는 $O(T)$, 즉 선형적으로 증가하므로 연산량 자체는 매우 효율적입니다.
  • 결론: 현재 구현상 학습 속도(Latency)가 Full Attention보다 느린 구간이 존재하며, 이는 향후 최적화(Custom Kernel 등)가 필요한 부분입니다.

 

 

Section 3.7 핵심 정리 노트 (Target: AI Researchers)

이 섹션은 TTT-E2E의 실용성을 **Inference(Prefill/Decode)**와 Training 두 가지 측면에서 분석하며, 특히 Training 단계에서의 Critical Limitation을 솔직하게 기술합니다.

  • Inference Efficiency (Prefill & Decode)
    • Standard Infrastructure Advantage:
      • Mamba 2, Gated DeltaNet 등 타 RNN 계열은 Efficient Memory I/O를 위해 Custom Kernel 작성이 필수적이거나, Hidden State를 GPU On-chip memory(SRAM)에 우겨넣어야 해서 State Size에 제약이 있음 (e.g., TTT-KVB는 LoRA 사용).
      • 반면, TTT-E2E는 Standard MLP Layer를 State로 사용하므로, 별도의 Custom Kernel 없이도 **Standard Parallelism Tools(Sharding)**를 사용하여 GPU 확장이 용이함.
    • Decode Strategy:
      • 매 토큰마다 TTT를 수행하지 않고, Mini-batch($1K$)가 찰 때마다 한 번씩 Update 수행.
      • Latency = (SWA Decode Latency) + (Periodic TTT Prefill Latency).
  • Training Efficiency (The Main Limitation)
    • Mechanism: Gradients of Gradients 연산 필요.
    • Performance Gap:
      • Long Context ($128K$): Full Attention 대비 1.2$\times$ Faster (Linear Scaling 덕분).
      • Short Context ($8K$): Full Attention 대비 3.4$\times$ Slower.
    • Bottleneck Analysis:
      • No FlashAttention: 현재 cuDNN FlashAttention은 Gradients of Gradients를 지원하지 않음.
      • Checkpointing Overhead: FLOPs는 토큰 당 상수($O(1)$)로 일정하지만, Latency는 Context Length $T$에 대해 $\log(T)$로 증가함 (Gradient Checkpointing through time 때문).
    • Practical Issue: 대부분의 Pre-training은 Short Context($8K$)에서 수행되므로, 현재 구현체는 학습 비용이 매우 비쌈.
  • Future Directions for Training Speed
    1. Custom Kernel: Gradients of Gradients를 지원하는 Custom Attention Kernel 개발 시 하드웨어 활용도(Utilization) 대폭 향상 가능.
    2. Initialization Strategy: TTT 없이 Standard Transformer로 Pre-training을 마친 후, TTT-E2E로 변환하여 Fine-tuning (RNN 연구들의 일반적 기법). 이렇게 하면 비싼 TTT 학습 구간을 최소화할 수 있음.

쉬운 설명 : "실전엔 강하지만, 수련 과정이 고되다"

1. 실전 투입 (Inference): "장비 탓을 하지 않는다"

  • 다른 최신 모델들(Mamba 등): "나를 빠르게 돌리려면 특수 제작된 칩 설정(Custom Kernel)이 필요하고, 메모리도 아주 정교하게 써야 해." (까다로움)
  • TTT 모델: "나는 그냥 흔한 부품(MLP)으로 만들어졌어. 기존에 쓰던 GPU 장비 그대로 써도 아주 빠르고, 확장하기도 쉬워." (범용성 좋음)

2. 수련 과정 (Training): "기초 훈련이 너무 느리다"

  • 문제점: 이 모델은 "학습하는 법을 학습"해야 합니다(미분을 두 번 함). 그래서 기초 체력을 기르는 짧은 글 학습(8K) 단계에서는 일반 모델보다 3.4배나 느립니다.
  • 희망: 하지만 글이 엄청나게 길어지면(128K), 일반 모델은 연산량이 폭발해서 느려지지만 TTT는 속도가 일정해서 오히려 1.2배 더 빨라집니다.
  • 결론: "나중에 가면 빠른데, 초반 레벨업 구간이 너무 힘들다."

3. 해결책은?

  • 저자들은 "일반 모델로 빨리 만렙(Pre-training) 찍고, 전직(Fine-tuning)해서 TTT 기술만 나중에 배우게 하면 되지 않을까?"라고 제안합니다.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

더보기

스크래치부터 학습

 

W가 있으면

W한테 학습 시킬 0~1000단어를 맞추게해서 loss를 측정하고 그걸 가지고
W1로 업데이트

w1로 1001~2000으로 loss 재고
w2 로 업데이트

w2fh 2001~3000으로 loss 재고

3개의 loss를 가지고 W를 업데이트 W부터 잘했다면 loss가 낮을텐데, 라는 방향으로 W가 업데이트 되고

 

이후 주어진 컨텍스트를 잘 기억할 수 있게 mlp를 학습하게 됨.

 

학습하는 mlp는 1/4 정도의 끝단 트랜스포머 블록의 mlp를 학습.

 

 

 

아쉬운점 

NTP를 확장하지 않고, TTT학습에서 지식 인출을 학습했으면 어땠을까 하는 아쉬움이; 남음. 그렇다면 바늘 찾기 벤치마크 성적도 매우 올랐을 것으로 예상.

 

 

 

별점

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

별점 4점 / 5점

단순하게 기억하기 위한 파라미터를 사용하는 것 이 아닌, 다른 모든 파라미터들을 기억하기 위한 파라미터를 위해 학습되는 것이 흥미로움, AGI로 향하는 기억력이라는 부분을 해결 할 수 있는 좋은 아이디어라고 생각, 하지만 학습 비용이 비싸고 전체 파라미터 중에 극히 일부분을 업데이트 하는 것으로는 많은 정보를 담을 수 없음. 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

'논문리뷰' 카테고리의 다른 글

VLM : 빠른 논문 리뷰 : Generate, but Verify: Reducing Hallucination in Vision-Language Models with Retrospective Resampling  (0) 2026.01.22
LLM : 빠른 논문 리뷰 : Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models  (0) 2026.01.21
LLM : 빠른 논문 리뷰 : WHAT MAKES MATH PROBLEMS HARD FOR REINFORCEMENT LEARNING: A CASE STUDY  (0) 2026.01.16
U-net : 논문 리뷰 : nnU-Net Revisited: A Call for Rigorous Validation in 3D Medical Image Segmentation  (1) 2026.01.15
Diffusion : 논문 리뷰 : GENERALIZATION IN DIFFUSION MODELS ARISES FROM GEOMETRY-ADAPTIVE HARMONIC REPRESENTATIONS  (0) 2026.01.15
'논문리뷰' Related Articles more