ATLAS: Learning to Optimally Memorize the Context at Test Time

Behrouz, A., Li, Z., Kacham, P., Daliri, M., Deng, Y., Zhong, P., Razaviyayn, M., & Mirrokni, V. (2025). Atlas: Learning to optimally memorize the context at test time. arXiv. https://arxiv.org/abs/2505.23735

이전편을 읽고오자

4. DeepTransformers: Transformers with Deep Memory

변형한 메모리 용량과 Omega rule이 실제 Transformer에 어떤 영향을 미치는지를 설명

4.1 Online and Local Context Optimization of Memory

Sliding Window Attention과 연결

Softmax block은 Nadaraya-Waston estimator을 이용해서 non-parametric solution으로 re-formulated 될 수 있다.

• $M^* = arg min_M \displaystyle \sum^L_{i=1}s(k_i, q) ||v_i -M||^2_2 = {{\sum^L_{i=1} s(k_i, q)v_i}\over{\sum^L_{j=1}s(k_j,q)}}$

이를 다시 과거 c개의 토큰으로 제한하면

• $M^* = arg min_M \displaystyle \sum^t_{i=t-c+1}s(k_i, q) ||v_i -M||^2_2 = {{\sum^L_{i=t-c+1} s(k_i, q)v_i}\over{\sum^L_{j=t-c+1}s(k_j,q)}}$
• 이는 Sliding Window Attention과 동일함
• Attention은 non-parametric한 방법으로 global한 최적화를 수행
  • 최신 recurrent model은 online-learning 방식
    • non-parametric: 정보를 압축하지 않고 데이터를 그대로 활용
    • online learning: 새로운 토큰에 대해 직전 요약된 메모리만 참조
  • Sliding window RNN과 Omega rule은 과거 토큰의 일정한 context window에 대해 메모리를 최적화해서 context 자체를 기억하게 만듦

Deep Linear Attention

linear attention의 행렬 메모리를 MLP로 대체

• deep memory 모듈로 메모리를 graident descent로 최적화
• $M_t = \alpha_t M_{t-1} - \eta_t ∇l(M_{t-1}; \Phi(k_t),v_t)$
  • $l(M_{t-1}; \Phi(k_t), v_t) = <M_{t-1}(\Phi(k_t)),v_t>$
• chunk 단위 병렬화로 효율적으로 사용할 수 있음

Sliding Window Linear Attention

• Associative memory에서 linear attention에 따라 dot product similarity를 attnetional bias로 사용
• $l(M_t;k_i,v_i) = <M_t(k_i),v_i>$
• gradient descent로 최적화
  • $M_t = \alpha_t M_{t-1} - \eta_t ∇ \displaystyle \sum^t_{i=t-c+1} l(M_{t-1}; \Phi(k_i), v_i)$
  • $M_t = M_{t-1} + \displaystyle \sum^t_{i=t-c+1}\gamma^{(t)}_i v_i \Phi(k_i)^\top$
  • c = 1, $\Phi(x) = x$일때 기존 linear attention과 동일

4.2 Memory Capacity and Exponential Kernels

기존 Transformer의 softmax attention

• $y_i = {{1}\over{\sum^i_{l=1}exp(q_i^\top k_l / \sqrt{d_{in}}}}\displaystyle \sum^i_{j=1}exp(q^\top_i k_j / \sqrt{d_{in}})v_j$
  • 이때 exp는 분리되지 않기에 recurrent한 표현이 어려움
• exp는 $\Phi_p$와 비교했을 때 입력을 무한차원 공간으로 매핑하는 feature map으로 볼 수 있음
  • $\Phi^*(x) = \begin{pmatrix} 1 \\ {{x}\over{\sqrt 1}} \\ {{x^{\otimes2}}\over{\sqrt{2!}}} \\ . \\ .\end{pmatrix}$
  • $\Phi_p(x) = x^{\otimes p}$
    • $x^{\otimes p}$는 Kronecker 곱을 이용한 self-tensoring 연산
    • $exp(q_t^\top k_t) = \Phi^* (q_t)^\top \Phi^* (k_t)$
• 최종적으로
  $y_i = {{1}\over{\sum^i_{l=1} exp(q^\top_ik_l)}} \displaystyle \sum^i_{j=1} v_j \Phi^*(k_j)^\top \Phi^*(q_i)$
  $={{1}\over{\sum^i_{l=1} exp(q^\top_i k_l)}}(\displaystyle \sum^i_{j=1} \Phi^*(v_jk_j)^{\top})\Phi^*(q_i) = M_i \Phi^* (q_i)$

이를 활용해 Deep Linear Attention에서 $\Phi^*$를 활용해 DeepTransformers 정의

• $M_t = M_{t-1} - ∇<M_{t-1}(\Phi^*(k_t)),v_t>$
• 선형인 경우: $M_t = M_{t-1} - ∇<M_{t-1}\Phi^*(k_t), v_t> = M_{t-1} + v_t\Phi^*(k_t)^\top$
  $=\displaystyle \sum^t_{i=1}v_i\Phi^*(k_i)^\top$
• Output: $y_t = M_t\Phi^*(q_t) = \displaystyle\sum^t_{i=1} v_i exp(q_i^\top k_i)$

4.3 Deep Omega Tranformer(DoT)

위의 DeepTransformer에서 Hebbian Rule대신 Omega learning Rule 사용

• $M_t = M_{t-1} - ∇ \displaystyle \sum^t_{i = t-c+1} \gamma^{(t)}_i ||M(\Phi^*(k_i)) - v_i||^2_2$

Linear memory: unnormalized된 형태를 얻음

• $M_t = (I- \displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}_i \Phi^*(k_i) \Phi^*(k_i)^\top)M_{t-1} - \displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}_i v_i \Phi^*(k_i)^\top$

• 최종 출력: $y_t = M_t\Phi^*(q_t)$

$= (I-\displaystyle\sum^t_{i=t-c+1} \gamma^{(t)}_i \Phi^*(k_i) \Phi^*(k_i)^\top) M_{t-1}\Phi^*(q_t) - \displaystyle\sum^t_{i=t-c+1} \gamma^{(t)}_i v_i exp(q_t^\top k_i)$

Online Case: c = 1

• $M_t = (I-\eta_t \Phi^*(k_t)\Phi^*(k_t)^\top)M_{t-1} - \eta_t v_t \Phi^*(k_t)^\top$
• 출력: $y_t = M_t\Phi^*(q_t)$
  $= (I-\eta_t \Phi^*(k_t)exp(q_t^\top k_t)) M_{t-1} - \eta_t v_t exp(q_t^\top k_t)$
• feature map을 사용하기에 이전 메모리 상태를 이용해 value를 replace하는 기능도 가능함

5. ATLAS: A Locally Optima lMemory with High Capacity

Omega Rule

• 모델이 context를 기억하도록 만듦 + feature mapping의 메모리 용량 증가
• gradient descent에 의존 $\rightarrow$ local optimum에 머무르게 함

Muon optimizer

• 2차 정보를 근사 + 행렬 곱 연산으로 구성되어 병렬화 가능
• $M_t = \alpha_t M_{t-1} - \eta_t NewtonShulz-k(S_t)$
• $S_t = \theta_t S_{t-1} + ∇\displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}_i ||M(\Phi^*(k_i)) - v_i||^2_2$
  • c: local cantext 길이
  • k: Newton-Schulz 반복 횟수
• 성질: $k \rightarrow \infty$일때 $NewtonShulz-k(S_t)$는 $S_t$에 가장 가까운 semi-orthogonal 행렬로 수렴
  • 2차 정보를 더 정확하게 반영
  • 즉 k가 크면 더 정확한 메모리 업데이트, 작은 k는 더 빠른 계산

5.1 Parallel Training

Atlas의 병렬화 방법

• c = 1로 가정
• 시퀀스를 여러 개의 청크로 나눔
• 청크의 마지막 상태를 기준으로 gradient를 계산

Muon 제외, momentm만 사용

• $M_t = \alpha_t M_{t-1} + S_t$
• $S_t = \theta_t S_{t-1} - \eta_t ∇l(M_{t'}, k_t, v_t)$

gradient 정의 변경

• $u_t = ∇l(M_{t'}, k_t, v_t)$
• $M_t = \alpha_t M_{t-1} + S_t$
• $S_t = \theta_t S_{t-1} - \eta_t u_t$
• 전개: $S_t = \theta_t S_{t-1} - \eta_t u_t$

장점

• 모든 graident $u_t$를 병렬로 계산 가능

6. Experiment

Task

• Language Modeling: 일반적인 언어 모델링에서 컨텍스트가 길어져도 기존 모델보다 우수
  
• Common-Sense Reasoing: Transformer 수준 or 그 이상
• Recall-Intensive Tasks: 긴 문맥에서 정보를 찾아내는 과제: Atlas가 가장 강력
• Long Context Understanding: 매우 긴 시퀀스(수백만 토큰)에서 안정적
  • 특히 BABILong 벤치마크에서 10M context에 대해 +80% 정확도 향상
• Needle-in-a-Haystack: Atlas와 OmegaNet이 가장 잘찾음

Ablation Study

• Memory Architecture: Deeop memory 성능이 크게 향상
• Feature Mapping: 고차원 feature일수록 성능 차이가 큼
• Omega Rule: online, sliding window, global과 비교할 때 context 기반 학습이 가장 좋음
• Optimizer: Muon이 가장 안정적