사카나 AI, 트랜스포머용 범용 메모리 모델 공개…선택적 망각으로 효율 극대화

사람은 뇌로 들어오는 자극의 대부분을 적극적으로 잊어버림으로써 중요한 정보에 집중합니다. 하지만 트랜스포머 기반의 대규모 언어 모델은 긴 문맥을 처리할 때 모든 토큰을 무차별적으로 메모리에 쌓아두기 때문에 작업이 길어질수록 추론 비용이 눈덩이처럼 불어납니다.

일본 경제산업성 GENIAC 슈퍼컴퓨팅 지원을 받는 사카나 AI(Sakana AI)가 이 문제를 해결할 새로운 메모리 시스템 'NAMM(Neural Attention Memory Models)'을 공개했습니다. 연구 논문 'An Evolved Universal Transformer Memory'를 통해 인간의 선택적 망각을 모사한 메모리 관리 기법을 발표했습니다.

역전파 대신 진화 알고리즘으로 배운 망각

기존에도 H2O나 L2처럼 메모리 사용량을 줄이기 위한 시도는 있었지만, 사람이 직접 짠 규칙에 의존해 토큰을 잘라내다 보니 모델 고유의 성능이 깎이는 부작용이 컸습니다.

사카나 AI는 각 토큰을 기억할지 영구히 잊을지 결정하는 가벼운 신경망 분류기를 구축했습니다. 문제는 토큰을 버릴지 남길지 결정하는 동작이 미분 불가능한 이진 선택이라 기존의 역전파 경사하강법으로는 학습시킬 수 없다는 점이었습니다.

연구팀은 주특기인 진화 알고리즘을 해법으로 꺼냈습니다. 반복적인 변이와 선택 과정을 거쳐 오차 역전파 없이도 최적의 망각 결정을 내리는 네트워크 파라미터를 진화시켰습니다.

작동 과정은 각 토큰의 어텐션 값을 음향학에서 쓰이는 스펙트로그램 형태로 바꾼 뒤 지수이동평균으로 압축하고, 진화된 신경망이 이를 평가해 불필요한 토큰을 쳐내는 방식입니다.

텍스트로 익힌 감각을 영상과 로봇에 그대로 이식

이 시스템의 가장 독특한 점은 제로샷 전이 능력입니다. 트랜스포머의 공통 구성 요소인 어텐션 행렬만 바라보고 판단하기 때문에 모델의 아키텍처나 데이터 모달리티에 구애받지 않습니다.

Llama 3 8B 언어 모델 위에서 진화시킨 NAMM을 추가 훈련 없이 대형 모델인 Llama 70B에 붙이거나, 비디오 모델(Llava Next Video)에서 중복 프레임을 걸러내고, 강화학습 모델(Decision Transformer)에서 비효율적인 행동 토큰을 솎아내는 데 성공했습니다.

내부 동작을 뜯어보면 모델의 앞쪽 레이어에서는 지시문 같은 전체 맥락을 기억하다가, 뒤쪽 레이어로 갈수록 이미 반영된 지시문은 잊고 국소적인 세부 정보에 집중했습니다.

코딩 작업에서는 공백과 주석, 불필요한 템플릿 코드를 뭉텅이로 잘라냈고, 자연어에서는 문법적 군더더기를 털어내며 고유명사와 핵심 단어를 살렸습니다. 연구팀은 전체 학습 코드와 함께 36개 과제로 구성된 일본어 장문 벤치마크 ChouBun도 오픈소스로 공개했습니다.

왜 중요한가

긴 콘텍스트를 서빙하는 기업 입장에서 메모리 병목과 추론 속도는 직접적인 비용 문제입니다. NAMM은 기반 모델을 처음부터 다시 학습시키지 않고도 외부 모듈 형태로 붙여 메모리를 절감할 수 있는 실용적인 대안을 제시했습니다. 코드가 공개되어 장문 처리 파이프라인을 운영 중인 엔지니어라면 직접 테스트해 볼 가치가 충분합니다.

업계 관점에서는 무작정 KV 캐시 크기를 키우거나 아키텍처를 완전히 갈아엎지 않아도 효율을 끌어올릴 수 있음을 입증했습니다. 특히 텍스트에서 학습한 프루닝 전략이 비디오 프레임과 로봇 제어에도 그대로 작동한다는 사실은 트랜스포머 계열 모델 전반의 공통된 정보 처리 방식을 시사합니다.

미분 불가능한 토큰 삭제 문제를 진화 연산으로 우회하고 어텐션 행렬의 범용성을 짚어낸 점은 매우 영리한 접근으로 보입니다. 다만 한계도 명확합니다. NAMM의 추론 자체는 가볍지만 이를 찾아내는 진화 과정은 슈퍼컴퓨터의 대규모 연산 지원이 필요할 만큼 탐색 비용이 큽니다. 일반 개발팀이 자체 도메인 모델에 맞춰 처음부터 직접 진화시키기보다는, 사카나 AI가 공개한 가중치를 그대로 전이해 사용하는 형태가 현실적인 활용법이 될 것으로 보입니다.