최근 트렌드로 떠오른 '지속 학습(Continual learning)'을 효과적으로 수행할 방법이 등장했다. 특히 기업 에이전트는 긴 문서나 다양한 로그를 처리하며 새로운 지식을 계속 학습하는 데, 여기에서 발생하는 롱 컨텍스트 처리의 비용 폭증 문제를 해결할 수 있다는 설명이다.
엔비디아 연구진은 6일(현지시간) 언어 모델을 사전 학습한 뒤 실제 활용 과정에서도 지속적으로 학습하는 과정의 효율성을 극대화할 수 있는 ‘엔드투엔드 테스트 타임 트레이닝(TTT-E2E)’을 온라인 아카이브에 공개했다.
연구진에는 엔비디아의 연구 선임 디렉터로 최근 범용 게임 에이전트 ‘나이트로젠’의 개발에도 참여했던 최예진 스탠포드대학교 교수가 포함됐다.
이번 연구는 기존 대형언어모델(LLM)이 사전 학습 단계에서 방대한 데이터를 통해 사실과 패턴을 최대한 많이 암기한 뒤, 배포 이후에는 더 이상 학습하지 않고 가중치가 고정된 상태로 답변만 생성하는 한계를 넘기 위한 것이다.
TTT-E2E는 모델이 사전에 모든 정보를 기억하도록 하는 대신, 실제 서비스 과정에서 새로 입력되는 정보를 처리하면서 실시간으로 학습하고 적응하도록 설계됐다. "아키텍처 설계 문제가 아닌, 지속 학습 문제"라는 설명이다.
이 과정에서 AI는 문서나 대화를 처리하며 다음에 올 토큰을 예측하고, 그 과정에서 중요한 내용을 내부에 정리해 저장한다. 덕분에 긴 문서를 다루거나 복잡한 기록을 분석해야 하는 기업 환경에서도 속도는 유지하면서 이해력은 높일 수 있는 방법이 될 수 있다고 연구진은 설명했다.
TTT-E2E는 기존 트랜스포머 구조를 그대로 사용하지만, 모든 문맥을 한꺼번에 보는 ‘풀 어텐션’ 대신 '슬라이딩 윈도우 어텐션’을 적용한다. 즉, 최근에 나온 일부 문맥만 집중해서 처리하는 방식이다.
이렇게 계산량을 줄였음에도 불구하고, 아주 긴 문서를 다룰 때의 정확도가 풀 어텐션 트랜스포머와 거의 비슷하다는 것이다. 또 문맥이 아무리 길어져도 응답 속도가 느려지지 않는다. 추론 지연 시간이 문맥 길이에 따라 늘어나지 않아, 동작 효율은 RNN(순환신경망)과 비슷한 수준을 달성했다.
연구진에 따르면 1640억 토큰으로 학습한 30억(3B) 매개변수 모델 기준으로, TTT-E2E는 12만8000토큰 컨텍스트에서 풀 어텐션 트랜스포머 대비 약 2.7배 빠른 추론 속도를 기록했다. 이는 문맥이 길어질수록 계산 비용이 급증하는 기존 트랜스포머의 한계를 크게 완화한 결과다.
연구진은 기존의 언어 모델들이 긴 문맥에서 성능이 떨어지는 이유로 ‘압축 부재’를 꼽았다. 풀 어텐션 모델은 모든 토큰을 거의 그대로 기억할 수 있어 정확도는 높지만, 문맥이 길어질수록 계산 비용이 급격히 늘어난다. 반대로 선형 시간 모델은 계산은 빠르지만, 시간이 지나면 이전 정보를 제대로 유지하지 못해 장기 기억력이 약하다는 한계가 있다.
TTT-E2E는 이 두 방식 사이의 간극을 문맥 압축으로 메운다. 슬라이딩 윈도우에서 벗어난 정보들을 그냥 버리는 대신, 다음 토큰 예측 과정에서 중요한 내용만 추려 모델의 가중치에 압축해 저장한다. 이렇게 하면 문서 앞부분의 핵심 정보가 형태를 바꿔 장기 기억으로 남게 된다.
이를 가능하게 하기 위해 모델 내부에는 이중 구조가 적용된다. 고정(static) 레이어는 사전학습으로 얻은 일반 지식을 유지하고, 동적(dynamic) 레이어는 현재 읽고 있는 문서의 정보를 실시간으로 반영한다.
또 가중치 업데이트는 모델 전체가 아니라 일부 다층신경망(MLP) 블록에만 제한적으로 수행되기 때문에, 학습 중 불안정해질 위험도 줄였다. 연구진은 이런 설계가 효율성과 안정성을 동시에 확보하는 핵심이라고 설명한다.
TTT-E2E의 또 다른 핵심 특징은 메타러닝이다. 연구진은 모델이 배포된 이후에도 스스로 학습할 수 있도록, 훈련 단계에서부터 ‘추론 중 학습’을 미리 연습하는 방식을 택했다.
구체적으로 학습 과정은 두 단계로 나뉜다. 먼저 내부 루프에서는 모델이 실제 사용 상황처럼 텍스트를 연속적으로 읽으면서, 다음 토큰을 예측하는 과정에서 아주 작은 가중치 업데이트를 수행한다. 이는 추론 중에 모델이 어떻게 적응할지를 그대로 시뮬레이션한 것이다.
그다음 외부 루프에서는, 이런 적응 과정이 더 빠르고 정확하게 일어나도록 모델의 초기 가중치 자체를 조정한다. 다시 말해, 모델에게 지식을 더 많이 외우게 하는 것이 아니라, 새로운 정보를 잘 배우는 ‘출발 상태’를 만들어주는 것이다.
연구진은 이 메타러닝 구조 덕분에 TTT-E2E가 실제 환경에서도 안정적으로 학습하며 긴 문맥을 처리할 수 있다고 설명한다.
컨텍스트 길이에 따른 스케일링 결과: 테스트 손실(왼쪽)과 지연 시간 (사진=arXiv)
실험 결과도 인상적이다. 기존의 효율적 모델들은 문맥 길이가 약 3만2000토큰을 넘어서면 성능 향상이 멈추거나 오히려 떨어졌지만, TTT-E2E는 풀 어텐션 트랜스포머처럼 문맥이 길어질수록 성능이 계속 개선됐다. 일부 실험 조건에서는 풀 어텐션 모델보다 더 낮은 퍼플렉서티(더 좋은 성능)를 기록하기도 했다.
한계도 분명하다. 긴 문서 속에 숨겨진 암호처럼 임의적이고 고립된 정보를 정확히 찾아내는 ‘건초더미 속 바늘 찾기(Needle in a Haystack)’ 유형의 S-NIAH 테스트에서는, 모든 토큰을 거의 완벽하게 기억하는 풀 어텐션 방식이 압도적으로 우수했다. 이는 TTT-E2E가 세부 정보를 그대로 보존하기보다는, 핵심 내용과 맥락을 압축해 저장하는 데 초점을 맞춘 구조이기 때문이다.
연구진은 이 점 때문에 검색 증강 생성(RAG)이 사라지지는 않을 것이라고 강조했다. 대신, 역할이 재정의된다. TTT는 사람의 뇌처럼 일반적인 지식과 맥락을 업데이트하는 방식이고, RAG는 필요할 때 정확한 정보를 찾아보는 메모장에 가깝다는 설명이다. 결국 두 방식은 경쟁 관계가 아니라, 서로를 보완하며 함께 사용될 가능성이 크다는 것이 연구진의 결론이다.
앞으로 AI의 메모리 방식이 ‘모든 것을 정확히 기억하는 구조’에서 ‘중요한 내용만 압축해 오래 기억하는 구조’로 바뀔 것이라고 전망했다. 모델이 세부 내용까지 완벽하게 기억할 수 있는 문맥 길이는 약 12만8000토큰 수준으로 유지하되, TTT 구조를 활용하면 그 너머의 정보도 수십억 토큰 규모로 압축해 저장하고 활용할 수 있다는 설명이다.
연구진은 관련 코드를 깃허브에 공개했으며, 추론 단계에서는 기존 트랜스포머 인프라에서도 바로 적용할 수 있다고 밝혔다.