Lifelong In-Context Learning with Transformers Requires Parametric Forms of Attention

기술 요약: 트랜스포머를 이용한 평생 인컨텍스트 학습 (Lifelong In-Context Learning)

문제 제기
트랜스포머를 기반으로 하는 현대의 AI 시스템은 유한한 데이터셋에 대한 오프라인 학습에 의존한다. 이들은 인컨텍스트 학습(In-Context Learning, ICL)을 통해 "지능의 불꽃(sparks of intelligence)"을 보여주지만, 표준 소프트맥스 어텐션(softmax attention) 메커니즘은 평생 학습(lifelong learning)을 위한 근본적인 확장성 장벽에 직면해 있다. 소프트맥스 어텐션의 이차적(quadratic) 특성은 계속해서 증가하는 키-값(Key-Value, KV) 캐시 유지를 요구하며, 이는 추론 비용과 메모리 점유율이 컨텍스트 길이에 따라 선형적으로 증가하게 만든다. 이는 고정된 하드웨어 예산 하에서 에이전트가 임의로 긴 시퀀스(예: 인간의 일생에 해당하는 토큰들)를 처리하는 것을 불가능하게 한다. 현재의 희소 어텐션(sparse attention) 방식은 토큰을 제거(eviction)함으로써 이를 완화하려 하지만, 과거의 경험을 일반적인 지식이나 태스크 관련 충분 통계량(sufficient statistics)으로 병합하는 데 실패하며, 이는 전체 컨텍스트를 모델링해야 하는 작업에는 불충분하다.

방법론 및 관점
본 논문은 어텐션을 단순한 검색 메커니즘이 아니라, **테스트 타임 회귀(test-time regression)**를 수행하는 온라인 학습 알고리즘으로 재정의한다.

• 회귀로서의 어텐션: 저자들은 어텐션 메커니즘을 트랜스포머 자체에 의해 생성된 자기 지도 학습 세트를 기반으로 키($k$)를 값($v$)으로 매핑하는 추정치 $m_t$를 구축하는 과정으로 본다. 쿼리($q$)는 레이블이 없는 테스트 포인트 역할을 하며, 출력은 예측된 값이다.
• 비매개변수적(Nonparametric) vs 매개변수적(Parametric):
  • 비매개변수적 (Softmax): 표준 소프트맥스 어텐션은 나다라야-왓슨 커널 추정기(Nadaraya-Watson kernel estimator)로 식별된다. 이는 원시 키-값 쌍을 "매개변수"로 저장하므로 메모리 성장이 무제한적이다.
  • 매개변수적: 본 논문은 매개변수적 형태의 어텐션을 옹호하며, 여기서는 키와 값 사이의 관계가 유한 차원의 온라인 업데이트 가능한 신경망(매개변수 $\theta_t$)에 의해 학습된다. 이는 계속 증가하는 KV 캐시를 상수 크기의 표현으로 대체한다.
• 테스트 타임 트레이닝 (TTT): 제안된 솔루션은 TTT를 활용하며, 여기서 매개변수 모델 $\theta_t$는 들어오는 키-값 쌍에 대해 손실 함수(예: 델타 규칙, 헤브 법칙 또는 오메가 규칙)를 최소화하도록 추론 시점에 경사 하강법을 통해 업데이트된다. 이는 선형 어텐션, 상태 공간 모델(SSM), 그리고 패스트 웨이트 프로그래머(fast-weight programmers)를 포괄한다.

핵심 기여

1. 이론적 통합: 본 논문은 다양한 최신 아키텍처(선형 어텐션, SSM, 패스트 웨이트 메모리)를 매개변수적 테스트 타임 회귀라는 관점으로 통합한다. 저자들은 이것들이 단순히 트랜스포머의 대안이 아니라, 평생 학습을 위해 필수적인 진화라고 주장한다.
2. 한계의 공식화: 비매개변수적 추정기(소프트맥스 등)는 과거의 경험을 추상화하거나 압축할 수 없고 단지 원시 토큰을 회상할 뿐이기에, 고정된 하드웨어 환경에서 평생 학습을 지원할 수 없음을 명시적으로 입증한다.
3. 평생 ICL을 위한 요구사항: 저자들은 어텐션 메커니즘이 평생 학습을 지원하기 위해 필요한 네 가지 속성을 정의한다:
   • 표현력이 있는 인컨텍스트 메모리 (단순 토큰 회상이 아닌 개념의 이해).
   • 컨텍스트 길이에 독립적인 토큰당 추론 비용.
   • 병렬화 가능한 업데이트 (현대적인 GPU/TPU 활용을 위함).
   • 제거가 필요 없는 메모리 (과거를 유한한 표현으로 병합).
4. 열린 질문의 식별: 새로운 아키텍처를 제안하기보다, 본 논문은 현재의 격차를 세 가지 핵심 연구 방향으로 정리한다:
   • 목적 함수 설계: 목표가 키-값 쌍의 완벽한 암기가 아니라면, 올바른 온라인 목적 함수는 무엇인가? 저자들은 정규화(예: 게이팅, 가중치 감쇠)가 즉각적인 컨텍스트에 과적합되는 것을 방지하고 일반화를 장려하는 데 중요하다고 제안한다. 또한 멀티 헤드 선형 어텐션이 암묵적인 정규화 역할을 한다는 점을 강조한다.
   • 장기 호라이즌 신호: 온라인 손실 함수가 어떻게 시간에 따른 추세를 효율적으로 포착할 수 있는가? 저자들은 즉각적인 업데이트(이는 파괴적 망각을 초래함)를 비판하며, 배치 업데이트(슬라이딩 윈도우)와 보조 상태(모멘텀, 2차 최적화 도구)가 유망하지만 아직 충분히 탐구되지 않은 방향임을 지적한다.
   • 자기 지도 레이블 생성: 키와 값은 어떻게 생성되어야 하는가? 저자들은 선형 투영만으로 충분한지, 아니면 온라인 학습기가 고차원 구조를 포착하도록 강제하기 위해 비선형 특징 맵과 더 풍부한 자기 지도 학습 태스크(예: 토큰당 다수의 KV 쌍)가 필요한지에 대해 의문을 제기한다.

결과 및 주장
본 논문은 새로운 경험적 벤치마크나 SOTA 결과를 내는 특정 혁신 아키텍처를 제시하지 않는다. 대신, 그 "결과"는 이론적이고 분석적이다:

• 소프트맥스 어텐션이 특정 비매개변수적 커널 회귀와 수학적으로 동일함을 증명하여, 그 메모리 제한의 이유를 설명한다.
• 멀티 헤드를 가진 선형 어텐션이 정규화 역할을 하여 모델이 훈련 데이터에 완벽하게 적합되는 것을 방지(이는 일반화를 위해 바람직함)함을 입증한다.
• 현재의 매개변수적 방법들(선형 어텐션이나 SSM 등)이 유망한 후보이지만, 제한된 메모리 용량, 비효적인 업데이트 규칙, 그리고 최적화되지 않은 목적 함수 설계로 인해 아직 미흡하다고 주장한다.

의의
본 논문의 의의는 패러다임의 전환에 있다: 어텐션을 정적인 검색 도구에서 동적인 온라인 학습기로 보는 관점의 변화이다.

• 현재의 트랜스포머 백본에 외부 검색 도구를 추가하는 것만으로는 충분하지 않다고 도전하며, 에피소드적 흔적(episodic traces)을 재사용 가능한 내부 구조로 변환하기 위해서는 내부적인 매개변수적 메모리가 필요하다고 주장한다.
• 이 분야의 연구가 업데이트 효율성, 메모리 용량, 목적 함수 설계라는 구체적인 병목 현상을 해결하기 위해 자기 지도 학습, 강화 학습, 지속 학습(continual learning)의 융합을 이루어야 한다고 촉구한다.
• 궁극적으로, AI 에이전트가 긴 호라이즌(수개월 또는 수년) 동안 작동하려면, 일정한 메모리 점유율을 유지하면서 실행 시간(runtime) 중에 경험으로부터 학습하는 매개변수적 어텐션 메커니즘을 채택해야 한다고 상정한다.