Learning to Remember, Learn, and Forget in Attention-Based Models

기술 요약: 어텐션 기반 모델에서의 기억, 학습 및 망각에 관한 학습

문제 정의
트랜스포머의 인컨텍스트 러닝(In-Context Learning, ICL)은 온라인 연상 메모리로서 기능하며 복잡한 시퀀스 작업에서 높은 성능을 가능하게 한다. 그러나 게이트형 선형 어텐션 모델과 상태 공간 모델(SSM)은 표준 트랜스포머(KV 캐시로 인해 선형 및 이차 스케일링이 필요한 모델)의 계산 및 메모리 병목 현상을 해결하기 위해 고정된 크기의 메모리 상태에 의존하는 경우가 많다. 이러한 고정된 용량은 "안정성-가소성 딜레마(stability-plasticity dilemma)"를 생성한다. 즉, 모델은 중요한 이전 지식을 덮어쓰지 않으면서(파괴적 망각) 새로운 정보를 학습해야 하며, 반대로 오래된 데이터의 과도한 보유로 인해 새로운 정보를 통합하지 못하는 상황(파괴적 기억)을 방지해야 한다. 신경과학에서 유래한 기존의 메타가소성(metaplasticity) 방법론들은 명확한 작업 경계에 의존하거나, 메모리가 성장하거나, 미분 불가능하다는 점 때문에 트랜스포머에는 적합하지 않은 경우가 많다. 또한, 현재의 선형 어텐션 모델들은 시퀀스 길이가 증가함에 따라 제한된 메모리 용량과 간섭 문제에 시달리는 경우가 많다.

방법론: 팔림사(Palimpsa)
저자들은 **팔림사(Palimpsa)**를 제안하며, 이는 ICL을 베이지안 메타가소성(Bayesian metaplasticity)을 통해 해결 가능한 지속적 학습(continual learning) 문제로 재구성한 셀프 어텐션 모델이다. 핵심 통찰은 각 어텐션 상태의 가소성이 축적된 지식을 포착하는 사전 분포(prior distribution)에 근거한 "중요도 상태(importance state)"에 따라 동적으로 조정되어야 한다는 것이다.

1. 베이지안 정식화: 어텐션 헤드는 변분 자유 에너지(variational free energy) 목적 함수를 최소화하는 것을 목표로 하는 베이지안 추론 문제로 모델링된다. 메모리 상태 $S$에 대한 사후 분포는 재귀적으로 업데이트된다. 파괴적 기억(사전 분포가 지나치게 집중되는 현상)을 방지하기 위해, 저자들은 누적된 사전 가중치를 효과적으로 절단하는 **가중 사후 분포(weighted posterior)**를 도입한다. 이는 망각 게이트 $\alpha_t = \exp(-A d_t)$에 의해 제어되는 방식으로, 각 단계에서 축적된 사전 가중치의 일부를 버림으로써 달성된다.
2. 메타가소성 업데이트 규칙: 대각 공분산 행렬 근사를 적용한 변분 추론(Variational Inference, VI)을 통해, 저자들은 폐쇄형(closed-form) 업데이트 규칙을 도출한다. 메모리 상태 $\mu_t$와 중요도 행렬 $I_t$(각 상태의 역분산 또는 정밀도를 나타냄)는 동시에 업데이트된다:
   • $I_t = \alpha_t I_{t-1} + (1 - \alpha_t)I_{prior} + \beta_t \otimes k_t^2$
   • $\mu_t = \alpha_t \frac{I_{t-1}}{I_t} \odot \mu_{t-1} + \frac{1}{I_t} \odot [(\beta_t \odot v_t) \otimes k_t]$
     여기서 $\beta_t$는 입력 의존적 중요도 인자로 작용하며, $I_t$는 각 상태의 학습률을 동적으로 조절한다. 높은 중요도(낮은 분산)는 안정성(낮은 가소성)으로 이어지고, 낮은 중요도는 더 빠른 학습을 허용한다.
3. 아키텍처 변형: 저자들은 두 가지 변형을 구현한다:
   • Palimpsa-D: $\beta_t$를 위한 투영(projection)과 잔차 항(residual term)을 추가하여 게이트형 델타 규칙(Gated Delta Rule) 모델인 Deltanets를 대체할 수 있는 드롭인 교체 모델이다.
   • Palimpsa-M: 입력 게이트 $d_t$가 망각 게이트와 입력 게이트의 이중 역할을 수행하도록 설계된 Mamba2 기반 모델이다.
4. 이론적 통합: 본 논문은 다양한 게이트형 선형 어텐션 모델(Longhorn, Deltanets, Mesanet)이 이 공통된 베이지안 프레임워크 내에서 특정 아키텍처 선택 또는 사후 분포 근사로서 나타남을 입증한다. 특히, Mamba2는 망각률이 매우 높아 중요도 행렬이 일정하게 유지되는($I_t \approx I_{prior}$), 즉 메타가소성이 비활성화된 팔림사의 특수한 사례임을 보여준다.

주요 기여

• ICL에서의 메타가소성: 불확실성에 기반하여 상태의 가소성을 동적으로 조정함으로써, 인컨텍스트 파괴적 망각과 파괴적 기억을 완화하는 베이지안 메타가소성을 도입하였다.
• Palimpsa 모델: 효율적인 GPU 훈련을 위한 커스텀 커널로 구현되었으며, 오래된 지식을 점진적으로 방출함으로써 파괴적 기억을 방지하는 확장 가능한 어텐션 블록이다.
• 통합 프레임워크: 기존의 게이트형 선형 어텐션 및 상태 공간 모델을 하나의 공통된 베이지안 프레임워크로 통합하였으며, Mamba2가 팔림사의 극한(limit)임을 밝혀냈다.
• 변환 파이프라인: 비메타가소성 모델(예: 사전 학습된 Mamba2)을 메타가소성 모델로 변환하는 방법으로, 비메타가소성 커널의 속도를 활용해 사전 학습을 수행하고 메타가소성의 메모리 용량을 활용해 미세 조정을 수행하는 과정을 제시한다.

실험 결과
저자들은 170M, 760M, 2.7B 파라미터 모델을 사용하여 합성 데이터, 상식 추론, 롱 컨텍스트 벤치마크에서 팔림사를 평가하였다.

• 합성 데이터 (MQAR): Multi-Query Associative Recall 벤치마크에서 팔림사는 비메타가소성 베이스라인(Gated Deltanet 및 Mamba2 포함)을 일관되게 능가한다. 성능 격차는 시퀀스 길이가 길어질수록 벌어지며, 이는 메타가소성이 복잡한 암기 작업에서 덮어쓰기 문제를 효과적으로 완화함을 보여준다.
• 언어 모델링 및 상식 추론: FineWeb-Edu 데이터셋에서 메타가소성 변형(Palimpsa-D 및 Palimpsa-M)은 비메타가소성 대응 모델 및 표준 트랜스포머에 비해 LAMBADA, PIQA, WinoGrande와 같은 작업에서 우수한 퍼플렉시티(perplexity)와 정확도를 달성했다. 구체적으로, 760M 규모에서 Palimpsa-D는 Gated Deltanet 베이스라인보다 평균 정확도에서 약 0.8포인트 앞선다.
• 롱 컨텍스트 성능: 2.7B 파라미터 모델 실험에서, 가장 약한 망각(가장 큰 유효 메모리 창)을 가진 8개의 Mamba2 레이어를 팔림사 블록으로 정밀하게 교체한 결과, 상식 추론에서 +0.84 포인트, LongBench 벤치마크에서 +1.1 포인트의 개선을 보였다. 메타가소성 모델은 짧은 컨텍스트 검색 작업에서는 약간 낮은 성능을 보였으나, 컨텍스트 길이가 길어짐에 따라 성능 저하가 훨씬 느리게 나타나 32k 토큰 지점에서 베이스라인을 추월했다.
• 학습 역학: 분석 결과, 메타가소성 비율(가소적 시냅스와 공고화된 시냅스의 차이)이 훈련 중에 증가하며, 이를 통해 모델이 새로운 데이터에 대한 가소성을 유지하면서도 중요한 정보를 안정화할 수 있음을 보여준다.

의의 및 주장
본 논문은 팔심사가 신경과학(메타가소성)과 베이지안 추론의 개념을 통합함으로써 고정 크기 메모리 모델의 안정성-가소성 딜레마에 대한 원칙적인 해결책을 제공한다고 주장한다. 그 의의는 다음과 같다:

1. 이론적 통찰: 서로 다른 선형 어텐션 아키텍처를 단일 베이지안 프레임워크로 통합하고, Mamba2를 망각이 지배적인 특정 극한 모델으로 식별하였다.
2. 실용적 확장성: 사전 학습(빠른 비메타가소성 커널 사용)과 미세 조정(메타가소성 커널을 사용하여 메모리 용량 확장)을 분리함으로써 대규모 메타가소성 모델의 사용을 가능하게 한다.
3. 성능: 메타가소성이 표준 트랜스포머의 이차적 메모리 비용을 요구하지 않으면서도 롱 컨텍스트 시나리오에서 메모리 보유 및 추론 능력을 일관되게 향상시킨다는 것을 입증하였다.

저자들은 이 접근 방식이 엄격한 메모리 제한 내에서 정보가 학습되고, 기억되고, 잊히는 방식을 최적화함으로써 자원이 제한된 시스템(엣지 디바이스)에서 고성능 시퀀스 모델링을 배포하는 데 기여한다고 결론짓는다.