KV 바인딩 기반 테스트 시간 학습은 사실상 선형 어텐션이다

테스트 시간 학습(TTT)이 데이터를 기억하는 '메모리' 방식이 아니라 '선형 어텐션'의 변형임을 수학적으로 증명했다. 이를 통해 복잡한 TTT 구조를 단순화하고 연산 속도를 4배 이상 높일 수 있는 병렬 처리의 길을 열었다는 점에서 큰 의의가 있다.

기존 TTT는 모델이 새로운 데이터를 볼 때마다 내부 가중치를 업데이트하여 해당 정보를 '기억'한다고 믿어왔다. 하지만 실제로는 더 잘 기억할수록(Loss가 낮을수록) 성능이 떨어지는 기현상이 발생했다. 이는 TTT가 정보를 저장하는 메모리가 아니라, 입력값들 사이의 관계를 계산하는 어텐션 메커니즘으로 작동하기 때문이다.

연구팀은 TTT의 내부 업데이트 과정을 수식으로 풀어내어, 이것이 사실은 입력값들 사이의 관계를 계산하는 '선형 어텐션'과 동일하게 작동함을 발견했다. 즉, 가중치를 바꾸는 행위가 정보를 저장하는 게 아니라 어텐션의 쿼리, 키, 값을 생성하는 과정이었던 것이다.

이 발견으로 인해 TTT는 더 이상 순차적인 업데이트에 얽매일 필요가 없어졌다. 선형 어텐션처럼 한꺼번에 계산(병렬화)할 수 있게 되어 속도가 비약적으로 향상되었으며, 복잡한 최적화 기법 없이도 높은 성능을 유지할 수 있게 되었다.

TTT-KVB의 내부 루프 업데이트 과정을 선형 어텐션 연산자로 변환하는 이론적 틀을 제시한다. 입력 토큰 x가 내부 MLP를 통과하여 쿼리, 키, 값을 생성하고 그래디언트 업데이트를 거치는 과정이 선형 어텐션의 상태 행렬(State Matrix) 업데이트와 수학적으로 일치함을 보여준다.

모멘텀이 포함된 SGD 업데이트 환경에서도 TTT가 선형 어텐션의 가중 합산 형태로 표현될 수 있음을 증명한다. [이전 그래디언트와 현재 그래디언트 → 모멘텀 계수 적용 → 누적 업데이트 → 가중치가 적용된 값(Value) 벡터 생성] 순으로 계산되어 최종적으로 선형 어텐션의 구조를 띤다.

가중치 정규화(Weight Normalization)를 제거하여 상태 업데이트의 결합 법칙(Associativity)을 회복한다. [상태 A + 상태 B → 결합 가능 → 병렬 프리픽스 스캔 적용] 과정을 통해 순차적 계산을 병렬 계산으로 전환하여 연산 효율성을 확보한다.

LaCT-LLM 및 ViTTT 모델에서 실험한 결과, 내부 루프의 그래디언트 하강을 상승으로 바꾸어도 성능이 유지되거나 오히려 향상됨을 확인했다. LaCT-LLM의 Perplexity 수치는 Baseline 16.43에서 Gradient Ascent 적용 시 16.19로 개선되는 결과를 보였다.

제안된 병렬화 기법을 적용했을 때, 기존 순차적 방식 대비 추론 처리량이 최대 4.0배 향상되었다. 또한 복잡한 구성 요소를 제거한 단순화된 모델(Variant 1)이 언어 모델링(Perplexity 15.93)과 이미지 분류(Top-1 Acc 79.63%) 등 모든 태스크에서 가장 우수한 성능을 기록했다.

학습 속도 측면에서도 병렬 구현을 통해 기존 대비 1.19배의 엔드투엔드 속도 향상을 달성하면서도 모델의 수렴 성능은 동일하게 유지함을 입증했다.

TTT-KVB의 최종 레이어가 선형적이고 편향(bias)이 없을 때, 한 단계의 그래디언트 업데이트 후의 출력은 o = q(S0 + kᵀv) 형태의 선형 어텐션으로 귀결됨을 수학적으로 유도했다.

가중치 정규화(Weight Normalization)가 병렬화를 방해하는 핵심 요소임을 식별했다. 정규화는 상태 업데이트의 결합 법칙을 깨뜨리기 때문에, 이를 제거함으로써 병렬 프리픽스 스캔(Parallel Prefix Scan) 적용이 가능해졌음을 이론적으로 분석했다.

기존 TTT 모델에서 사용되던 토큰별 학습률(Per-token learning rate)이나 복잡한 모멘텀 계수가 선형 어텐션의 값(Value) 벡터에 흡수될 수 있는 중복된 요소임을 밝혀내어 아키텍처 최적화의 근거를 마련했다.

본 연구의 분석은 내부 루프의 최종 레이어가 선형적이고 편향이 없는 경우로 제한된다. 비선형 최종 레이어를 가진 모델이나 더 복잡한 어텐션 메커니즘과의 완전한 통합은 향후 연구 과제로 남겨져 있다.