LLM - Google, ATLAS

 

 

Figure 1: Comparison of learning to memorize (Left) individual tokens, and (Right) the context., https://arxiv.org/abs/2505.23735

 

 

 

 

 

Figure 3: Visualization of the Atlas’s (and our other variants’) architecture, and its hybrid counterpart with SWA.

 

‘트랜스포머’ 아키텍처는 엄청난 성과를 보여주는 놀라운 아키텍처지만, 기본적으로 ‘모든 토큰 쌍 (Token Pair)을 살펴보는 어텐션 메커니즘의 속성 때문에, 그 훈련 비용이 사이즈와 함께 기하급수적으로 증가합니다. 물론, RetNet, RWKV, Mamba, 심지어 구글이 자체적으로 만든 Titans 같은 많은 대안들이 상태공간 기법 등의 다양한 방법으로 이 문제를 해결하고자 시도하고 있습니다. 이런 대안들은 더 빠르게 더 적은 비용으로 처리를 하게 해 주기는 하지만, ‘정말 긴 시퀀스에서 정보를 진정으로 이해하고 기억하는데 있어서는 여전히 많은 모델들이 벽에 부딪히고’ 있는 것도 사실입니다.

 

살짝 기술적으로 이해를 해 보자면, 제가 보기에는, 핵심적인 아이디어는 모델의 메모리를 단순히 정보를 집어넣는 수동적 저장소로 보기보다는 특히 추론 과정에서 능동적으로 최적화를 실행하는 구성 요소로 보고 다루는 겁니다. 다시 말해서, ‘메모리가 마지막에 본 토큰에만 반응하는게 아니라, 메모리가 최근 토큰들의 윈도우를 되돌아보면서 뭐가 중요한지, 어떻게 보관할지를 스스로 지능적으로 결정하는 방법’을 제안하고 있어요.

 

Image Credit: “Atlas: Learning to Optimally Memorize the Context at Test Time” 논문

 

여기서 ‘오메가 규칙’이라는 게 등장합니다. 오메가 규칙이 뭘까요?

 

대부분의 순환 신경망 모델은 ‘현재의 입력’을 바탕으로 해서 메모리를 업데이트하는데, 오메가 규칙은, “잠깐만, 최근 'c'개의 토큰들(예를 들어 최근 50개)을 살펴보고 이것들을 모두 함께 활용해서 메모리 상태를 최적화하자"고 하는 겁니다. 이렇게 해서, 모델이 단순히 ‘고립된 사실들’이 아닌 ‘맥락’을 기억하는 법을 배울 수 있다는 겁니다 - 즉, 마지막에 일어난 독립적인 일 뿐 아니라 최근의 더 넓은 경험을 바탕으로 이해를 업데이트하는 개념이죠.

 

논문을 보면, 이런 접근법이 특히 BABILong 같이 ‘아주 긴 문서에 흩어져 있는 사실을 바탕으로 언어 모델이 얼마나 장기 추론이나 정보 검색을 잘 하는지 평가하는 테스트’에서 실제로 좋은 성과를 보여주는데, 무려 1천만 토큰까지 이어지는 시퀀스에서 강력한 성능을 보여준다고 해요.

 

Image Credit: “Atlas: Learning to Optimally Memorize the Context at Test Time” 논문

 

그런데, Atlas가 이런 하나의 핵심 아이디어로만 구성된 건 아닙니다 - 여러 개의 보완적인 도구 셋을 소개하고 있어요:

 

1. 메모리 용량의 향상

 

Atlas는 키(Key)와 쿼리(Query)에 다항식 및 지수 특성 매핑을 적용해서, ‘운용할 수 있는’ 메모리의 용량을 증가시킵니다. 말이 좀 어려운데, 실제 메모리 저장소는 원래 크기 그대로 유지하면서, 정보를 찾고 비교할 때만 키와 쿼리를 더 복잡한 형태(다항식, 지수 등)로 임시 변환합니다. 마치 서랍장 개수는 그대로 두고 물건을 찾을 때만 더 정교한 인덱스 시스템을 사용하는 것처럼, 핵심적인 모델의 파라미터를 늘리지 않고도 더 많은 패턴을 구별하고 기억할 수 있게 만들어 줍니다.

이 접근법은 PolySketchFormer와 Hopfield 네트워크 확장을 포함한 이전 여러 가지 연구의 아이디어를 기반으로 만들어졌다고 하구요.

 

2. 더 효과적인 메모리의 조정

 

메모리를 업데이트할 때, 기본적인 경사하강법을 사용하는 대신, Atlas는 2차 정보를 활용하는 Muon 최적화기를 통합합니다. 이것도 조금 풀어서 말씀드리자면, 기본적 경사하강법으로 단순히 새 정보를 더하거나 평균내는게 아니라, Muon이라는 고급 최적화기를 사용해서 "어떻게 하면 이 메모리가 가장 유용하게 쓰일까?"라는 최적화 문제를 실시간으로 푸는 겁니다 - 마치 산을 오를 때, ‘발 밑의 기울기’만 보는게 아니라 ‘주변 지형의 곡률’까지 파악하면 더 나은 최적 경로를 찾을 수 있는, 즉 1차 정보(기울기)와 2차 정보(곡률)를 모두 활용해서 메모리가 더 안정적이고 효과적으로 중요한 맥락을 기억하도록 해 주는 기법입니다.

 

이 기법을 사용하면, 메모리 모듈을 더 효과적으로 업데이트해서 일반적인 최적화에 따르는 함정을 피하는 데 도움이 됩니다. 이 아이디어는 훈련 중에 더 고급의 최적화기를 사용하는 것과 비슷하지만, 여기서는 ‘추론 과정에서의 메모리 메커니즘’에 적용하는 거죠.

 

3. 트랜스포머 아키텍처의 재검토

 

더불어, 이 논문에서는 ‘고정된 어텐션’을 ‘오메가 규칙으로 관리하는 학습 가능한 메모리 모듈’로 대체해서 트랜스포머 아키텍처를 확장한 DeepTransformers와 DoT(Deep Omega Transformers)를 제안합니다. 이런 새로운 모델들을 오히려 ‘오리지널 트랜스포머의 일반화된 형태’로 바라본다면, 가장 기본적인 어텐션은 더 넓은 메모리 구성 방식의 한 가지 경우라고 해석할 수도 있을 겁니다.

 

이 논문은 상당히 ‘밀도가 높은’ 논문입니다. 연관 메모리 (Associative Memory) 이론, 최적화, 아키텍처 설계 등 다양한 영역에서 여러 가지의 새로운 실마리를 잘 엮어낸, 일종의 ‘작품’이라고 생가합니다. 전통적인 RNN 아키텍처의 ‘온라인 업데이트’ 특성, 즉 토큰이 하나씩 들어올 때마다 미리 정해진 규칙으로 순차적으로 메모리를 업데이트하는 방식에 의문을 제기해 왔던 여러 가지 연구의 계보를 따르고 있기도 합니다 - 즉, 나중에 들어온 맥락이 이전 정보의 중요도를 바꿀 수 있음에도 불구하고, 이미 업데이트된 메모리를 되돌아 재조정할 수 없다는 근본적인 한계에 도전한다는 것이죠.

 

‘온라인’ 방식과 ‘Atlas’ 방식의 차이

 

Atlas를 비롯한 최근의 연구들은, 이렇게 ‘순서대로 차례차례 처리’하는 기법을 포기하고, 대신 일정한 범위의 토큰들을 함께 보면서 "어떻게 기억하는 게 가장 유용할까?"라는 최적화 문제로 메모리 업데이트를 재정의하고 있는데, 이건 일기를 시간순으로 나열하는 것에서 하루를 돌아보며 핵심 순간들 중심으로 정리하는 것으로 패러다임이 바뀐 것과 같다고 볼 수 있습니다. Titans (역시 구글에서 나온), TTT (Test-Time Training) 같은 논문들도, 다이나믹하게 추론 시점에서 적응한다는 비슷한 주제를 탐구한 바 있구요.

 

물론, Atlas가 최종적인 답이 아닐 확률이 크죠. 진짜 시험대는, 이런 아이디어들이 얼마나 잘 스케일링되는지, 다른 사람들이 얼마나 쉽게 구현하고 발전시킬 수 있는지, 그리고 더 광범위한 작업에서도 높은 성능이 유지되는지 등이 관건이 될 겁니다 - 많은 애플리케이션들은 여전히 짧은 맥락의 윈도우 안에서도 큰 문제없이 작동하고 있으니까요.

 

그렇지만, 분명 Atlas가 제시하는 방향은 의미가 있습니다.

 

기존의 패러다임을 단순히 확장하는 것이 아니라, 더 똑똑한 시스템을 구축한다는 철학이니까요. 특히 방대한 정보를 다루어야 할 때 내부 메모리를 관리하고 최적화하는 방법을 능동적으로 학습하는 모델, 멋지지 않나요? 이 방향의 진정한 발전이 계속된다면, ‘다이나믹하게 맥락을 인식하면서 지속적으로 최적화하는 메모리’는 대규모의 데이터셋을 단순히 ‘처리’하는게 아니라 ‘제대로 이해하고 활용’하는 AI를 향한 중요한 한 걸음이 될 겁니다.

 

이제, 단순하게, 그리고 기계적으로 어텐션을 확장하는 시대는 지나고 있습니다. AI가 어떻게 학습하고, 기억하고, 추론해야 하는지에 대해 더 집중해서 세밀하게 탐구하는 시대가 오고 있습니다.

 

 

 

https://turingpost.co.kr/

https://arxiv.org/pdf/2505.23735