Titans 논문 쉽게 이해하기 - 테스트 시점에 기억을 학습하는 신경 메모리 구조

핵심 메시지

• 긴 시퀀스 모델링에는 두 갈래가 있다. 트랜스포머는 전체 문맥을 병렬로 볼 수 있지만 KV 캐시가 계속 커져 긴 문맥 추론 비용이 비싸고, RNN은 고정 크기 메모리로 압축하지만 병렬화가 어렵고 처음 보는 데이터에서 압축이 무너진다.
• Titans의 핵심은 "테스트 시점(추론 중)에 기억을 학습하는" 신경 메모리 모듈이다. 학습 때 굳어버린 압축 능력에 의존하지 않고, 실제로 입력이 들어오는 순간 그 데이터에 맞춰 메모리를 갱신한다.
• 신경 메모리는 사실상 하나의 가중치 행렬(선형층 또는 MLP)이며, 입력을 복원하는 재구성 손실을 줄이도록 경사하강으로 그 자리에서 갱신된다. 손실이 클수록 "놀라움(surprise)"이 크다고 보고, 옵티마이저의 모멘텀처럼 과거의 놀라움을 함께 반영해 갱신을 부드럽게 만든다.
• 토큰 하나씩 갱신하면 느리므로, 논문은 청크 단위로 병렬화하는 알고리즘을 제안한다(예: 100만 토큰을 1,000 크기 청크로 나눠 1,000단계로 처리).
• 메모리를 문맥으로 붙이는 방식, 게이트로 결합하는 방식, 층 자체로 쓰는 방식 등 여러 변형이 가능하다. 벤치마크에서는 Mamba 등 기존 구조보다 평균 점수가 앞섰지만, 발표자는 벤치마크를 그대로 믿지 말라고 당부한다.

쉽게 이해하기

이 영상은 "Titans: Learning to Memorize at Test Time" 논문을 우마르 자밀이 직접 그림을 그려가며 해설한 강의다. 그는 논문을 문장 그대로 읽어주기보다, 배경지식부터 문제와 해법까지 스스로 이해할 수 있게 짚어주는 방식을 택한다.

먼저 시퀀스 모델링의 두 축을 정리한다. 트랜스포머는 어텐션으로 전체 입력을 한꺼번에 보고 각 위치를 병렬로 계산할 수 있어 GPU를 최대한 활용하지만, 다음 토큰을 예측할 때마다 지금까지의 모든 토큰(KV 캐시)을 들고 있어야 해 문맥이 길어질수록 추론 비용이 급격히 커진다. 반대로 RNN은 고정 크기의 은닉 상태(메모리)에 과거를 압축해 넣기 때문에 메모리가 커지지 않지만, 기본형은 병렬화가 안 되고 무엇을 남기고 무엇을 버릴지 스스로 정할 수 없다.

진짜 문제는 학습과 추론의 간극이다. 트랜스포머와 RNN을 섞은 하이브리드 구조는 학습할 때는 손실이 잘 떨어지지만, 학습 때 본 적 없는(분포 밖) 입력이 들어오면 RNN 계열 메모리가 압축을 제대로 못 하고, 그 뒤 어텐션도 필요한 정보를 끌어오지 못해 성능이 무너진다. Titans는 이 지점을 겨냥한다.

Titans는 추론 시점에도 학습을 멈추지 않는 별도의 신경 메모리를 둔다. 100만 토큰짜리 입력이 들어오면, 본 모델의 파라미터와 분리된 이 메모리를 "그 자리에서" 학습시켜 입력을 최대한 기억하게 만든 뒤, 필요한 정보를 꺼내 어텐션 층에 넘긴다. 논문은 이를 안쪽 루프(신경 메모리 학습)와 바깥 루프(본 모델 학습)로 부른다.

메모리 갱신의 원리는 간단하다. 입력을 키(key)와 값(value)으로 사영한 뒤, 메모리가 그 매핑을 복원하도록 재구성 손실을 정의하고 경사하강으로 가중치 행렬을 갱신한다. 이때 손실의 기울기 크기를 "놀라움"으로 해석하고, 미니배치 경사하강의 잡음을 다듬는 모멘텀처럼 과거의 놀라움을 더해 새로운 정보가 갑자기 사라지지 않도록 한다.

주요 인사이트

• "테스트 시점 학습(test-time training)"이라는 큰 흐름 안에서, 메모리 그 자체를 추론 중에 갱신한다는 발상이 이 논문의 참신함이다. 학습 때 굳은 압축기가 아니라, 지금 들어온 데이터에 맞춰 최적화되는 메모리를 쓴다.
• 신경 메모리가 결국 하나의 선형층(또는 MLP) 가중치라는 점이 중요하다. "기억한다"는 추상적 개념을 재구성 손실을 줄이는 경사하강이라는 구체적 연산으로 환원했다.
• "놀라움"과 모멘텀의 연결은 옵티마이저 이론과 자연스럽게 맞닿는다. 큰 손실=큰 놀라움으로 보고, 과거 기울기를 지수이동평균으로 섞어 급격한 변화를 억제한다.
• 토큰별 갱신의 속도 문제를 청크 병렬화로 완화한 점은 실용적 타협이다. 완전 병렬은 아니지만 GPU를 어느 정도 살릴 수 있다.
• 발표자는 지속 메모리(persistent memory) 토큰이나 세 가지 사용 방식(문맥/게이트/층)은 부차적이라고 본다. 핵심은 "어떻게 쓰느냐"가 아니라 "추론 중에 메모리를 학습한다"는 원리 그 자체다.

자주 묻는 질문

관련 AI 소식