In-Place 테스트 시점 학습: LLM을 위한 원활한 동적 적응 프레임워크

기존 LLM은 학습이 끝나면 지식이 고정되어 새로운 정보를 실시간으로 받아들이기 어렵지만, 이 논문은 모델 구조 변경 없이 추론 중에 가중치를 업데이트하는 방식을 제안한다. 이를 통해 4B 규모의 작은 모델로도 128k 이상의 초장문 문맥을 효과적으로 처리할 수 있는 성능을 입증했다.

기존의 Transformer 모델은 한 번 학습되면 추론 시에는 가중치가 고정된 상태로 작동한다. 이로 인해 모델이 처음 보는 긴 문서나 실시간으로 변하는 정보를 처리할 때, 오직 Attention Mechanism의 KV Cache에만 의존해야 하는 한계가 있다. 문맥이 길어질수록 메모리 사용량이 급증하고 과거 정보를 효율적으로 압축하여 기억하기 어려워진다.

In-Place TTT는 모델의 일부 가중치를 '동적 메모리'로 변환하여 이 문제를 해결한다. 구체적으로는 모든 Transformer 층에 존재하는 MLP 블록의 마지막 행렬을 추론 중에 미세하게 조정한다. 이는 마치 사람이 글을 읽으면서 중요한 내용을 머릿속에 실시간으로 요약하고 저장하는 것과 유사한 원리다. 특히 기존 TTT 연구들이 모델 구조를 통째로 바꿔야 했던 것과 달리, 이 방식은 이미 잘 학습된 모델의 구조를 그대로 유지하면서 성능만 업그레이드한다.

결과적으로 모델은 입력되는 텍스트를 읽어나가며 스스로를 최적화한다. 실험 결과, 4B 파라미터 모델이 이 기법을 통해 128k 토큰에 달하는 방대한 문맥에서도 정확도를 유지했으며, 이는 기존의 정적인 모델들이 긴 문맥에서 급격히 성능이 저하되던 현상을 획기적으로 개선한 결과다.

In-Place TTT는 기존 Gated MLP 구조를 활용한다. 입력 투사 행렬인 Wup과 Wgate는 고정된 '느린 가중치'로 두고, 최종 출력 투사 행렬인 Wdown만을 '빠른 가중치'로 설정하여 추론 중에 업데이트한다. [입력 활성화 값 Z → Wdown과 행렬 곱 → 출력 O] 과정을 거치며, 여기서 Wdown은 각 데이터 청크마다 갱신된다.

효율적인 연산을 위해 청크 단위 업데이트 전략을 사용한다. 전체 시퀀스를 일정한 크기(C=512~1024)의 청크로 나누고, 각 청크 내에서 한 번의 Gradient Descent 단계를 수행한다. [청크 내 입력 Z[i]와 목표값 V[i] 입력 → 손실 함수 계산 → Wdown 업데이트] 순으로 진행되며, 업데이트된 가중치는 다음 청크를 처리할 때 사용된다.

학습 목표는 Next-Token Prediction(NTP)과 정렬된 LM-Aligned Objective를 사용한다. 1D Convolution을 통해 미래 토큰의 정보를 포함하는 목표값 V_hat을 생성한다. [현재 토큰 임베딩 입력 → 1D Conv 및 투사 연산 → 미래 정보가 담긴 V_hat 출력] 과정을 통해, 가중치가 단순히 현재 토큰을 암기하는 것이 아니라 다음 토큰을 더 잘 예측할 수 있는 방향으로 정보를 압축하도록 유도한다.

RULER 벤치마크 평가 결과, Qwen3-4B 모델에 In-Place TTT를 적용했을 때 128k 문맥 길이에서 기본 모델 대비 월등한 정확도를 보였다. 특히 64k 이상의 긴 문맥에서 성능 차이가 두드러졌으며, 256k 길이로 확장했을 때도 뛰어난 일반화 성능을 유지했다. LLaMA-3.1-8B와 Qwen3-14B 모델에서도 일관된 성능 향상이 확인되어 범용성을 입증했다.

처음부터 학습(Pre-training from scratch)한 실험에서도 우수성이 증명됐다. 500M 및 1.5B 규모 모델에서 기존의 선형 어텐션(GLA)이나 다른 TTT 방식(DeltaNet, LaCT)보다 낮은 Perplexity를 기록했다. 이는 In-Place TTT가 기존 모델의 사후 강화뿐만 아니라 새로운 아키텍처 설계로서도 강력한 경쟁력을 가짐을 의미한다.

Ablation Study를 통해 설계의 유효성을 검증했다. 청크 크기가 512에서 1024 사이일 때 성능과 효율의 균형이 가장 좋았으며, LM-Aligned Objective에서 1D Convolution과 투사 행렬을 모두 사용할 때 가장 높은 성능을 기록했다. 또한 TTT 적용 레이어 수가 많아질수록(State size 증가) 성능이 비례하여 향상됨이 확인됐다.

In-Place TTT는 Transformer의 MLP 블록을 동적 연상 메모리(Associative Memory)로 재해석한다. Wdown 가중치는 키-값(Key-Value) 쌍을 저장하는 저장소 역할을 하며, 추론 시 각 청크마다 한 번의 델타 업데이트(ΔW)를 수행한다. 이 과정은 수학적으로 선형 어텐션의 업데이트 규칙과 유사하지만, 기존 가중치를 보존하면서 증분 업데이트를 수행한다는 점이 다르다.

구현 측면에서는 Context Parallelism(CP)과의 호환성을 위해 Associative Scan 알고리즘을 도입했다. 각 청크의 업데이트량을 독립적으로 계산한 후 Prefix Sum(누적 합)을 통해 전체 시퀀스에 대한 가중치 변화를 병렬로 적용한다. 이는 순차적인 업데이트와 수학적으로 동일한 결과를 내면서도 현대 하드웨어에서 연산 속도를 극대화한다.

이론적 분석을 통해 LM-Aligned Objective가 단순 Reconstruction보다 우월함을 증명했다. Induction Head 설정에서의 분석 결과, NTP 정렬 목적 함수는 정답 토큰의 Logit을 직접적으로 상승시키는 반면, 단순 복원 목적 함수는 예측 성능에 유의미한 영향을 주지 못함을 수식으로 입증했다. 이는 TTT가 언어 모델의 본질적인 작업에 기여하기 위해 목적 함수 설계가 핵심임을 시사한다.

논문은 추론 시 가중치 업데이트로 인한 계산 오버헤드가 존재함을 언급한다. 비록 병렬화를 통해 최적화했으나, 순수 정적 모델 대비 추가적인 연산 자원이 필요하다. 또한 매우 긴 문맥에서 가중치 업데이트가 누적될 때 발생할 수 있는 수치적 불안정성을 방지하기 위해 별도의 Clipping 메커니즘이 필요함을 명시했다.