🤖 machine learning

Mixture of Chapters: Scaling Learnt Memory in Transformers

이 논문은 트랜스포머 모델의 학습된 지식을 저장하고 조직화하기 위해 학습 가능한 희소 메모리 뱅크와 분량 기반 라우팅을 도입하여, 파라미터 수와 FLOP 수를 동일하게 유지하면서도 기존 모델보다 우수한 성능과 지식 유지 능력을 보여주는 새로운 확장 축을 제시합니다.

원저자: Tasmay Pankaj Tibrewal, Pritish Saha, Ankit Meda, Kunal Singh, Pradeep Moturi

게시일 2026-03-24

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Tasmay Pankaj Tibrewal, Pritish Saha, Ankit Meda, Kunal Singh, Pradeep Moturi

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

🧠 핵심 문제: AI 는 '공부'를 해도 '공부 노트'를 못 만듭니다

지금까지의 AI(트랜스포머) 는 책을 읽거나 데이터를 학습할 때, 그 지식을 **머릿속의 뇌세포 (모델 파라미터)**에 직접 새겨 넣습니다. 마치 우리가 시험을 볼 때 암기한 내용을 머릿속에서 끄집어내는 것과 비슷하죠.

하지만 이 방식에는 두 가지 큰 문제가 있습니다.

기억이 섞여요: 새로운 것을 배우면 예전 지식이 지워지거나 (망각), 서로 뒤섞여 버립니다.
찾기가 어려워요: 필요한 지식을 떠올리려면 뇌 전체를 뒤져야 하므로 계산 비용이 너무 많이 듭니다.

💡 이 논문이 제안한 해결책: "지식 도서관"과 "유능한 사서"

이 연구팀은 AI 에게 **별도의 '학습된 기억 은행 (Memory Bank)'**을 만들어 주었습니다. 그리고 그 기억을 효율적으로 찾기 위해 **'장 (Chapters)'**으로 나누고, **스마트한 '사서 (Router)'**를 배치했습니다.

이를 일상생활에 비유해 볼까요?

1. 기억 은행 (Learned Memory Bank) = 거대한 지식 도서관

기존 AI 는 지식을 뇌세포에 흩뿌려 두지만, 이 새로운 AI 는 별도의 거대한 도서관을 갖습니다.

이 도서관에는 **수십만 개의 '잠재 토큰 (Latent Tokens)'**이라는 책들이 있습니다.
이 책들은 AI 가 훈련하는 동안 스스로 내용을 채워 넣습니다.
AI 가 질문을 하면, 이 도서관에서 관련 있는 책들을 찾아 답을 얻습니다.

2. 혼합된 장 (Mixture of Chapters) = 도서관을 구역별로 나누기

도서관이 너무 크면 (26 만 권!), 모든 책을 한 번에 뒤지는 건 불가능합니다. 그래서 도서관을 **4,000 개의 '장 (Chapters)'**으로 나눕니다.

예: 1 장은 '과학', 2 장은 '역사', 3 장은 '요리'처럼 구분된 것이지요.
하지만 모든 장을 다 볼 필요는 없습니다.

3. 사서 (Router) = 똑똑한 안내원

질문이 들어오면, AI 는 유능한 사서에게 "이 질문과 관련된 책이 어디에 있을까?"라고 묻습니다.

사서는 전체 도서관을 다 뒤지는 대신, 가장 관련 있는 64 개의 '장'만 골라냅니다.
이 방식 덕분에 AI 는 거대한 도서관을 가지고 있으면서도, 계산하는 속도는 여전히 빠릅니다. (비유하자면, 26 만 권의 책이 있어도 매번 64 권만 보면 되니까요!)

🚀 이 방식이 왜 대단한가요?

이 연구팀은 이 방식을 적용한 AI 를 기존 AI 와 비교 실험했습니다. 결과는 놀라웠습니다.

1. 더 적은 비용으로 더 똑똑해짐 (Iso-FLOP)

같은 양의 전기를 쓰고 (계산 비용 동일), 같은 시간 동안 공부시켰을 때, '기억 은행'이 있는 AI 가 훨씬 더 낮은 실수율을 보였습니다.
마치 같은 시간 공부해도, 공부 노트를 잘 정리한 학생이 더 좋은 성적을 내는 것과 같습니다.

2. 잊어버리지 않음 (Retention)

기존 AI 는 새로운 것을 배우면 (예: 대화 능력을 가르침) 예전에 배운 사실 (예: 역사 지식) 을 잊어버리는 경우가 많습니다. 이를 '망각'이라고 합니다.
하지만 '기억 은행'이 있는 AI 는 새로운 것을 배우면서도 예전 지식을 도서관에 안전하게 보관해 두었습니다.
비유: 기존 AI 는 새 옷을 입으면 예전 옷을 버리는 반면, 이 AI 는 새 옷을 입으면서도 옷장 (기억 은행) 에 예전 옷을 깔끔하게 정리해 둡니다.

3. 학습 중에도 기억을 유지함

실험 결과, 학습을 계속해도 기억 은행을 '동결 (Frozen)'해 두기만 해도 성능이 유지되었습니다. 이는 AI 가 이미 배운 지식을 별도의 공간에 안전하게 저장하고 있다는 뜻입니다.

📝 한 줄 요약

"AI 가 머릿속에 모든 지식을 다 기억하려다 지치고 잊어버리는 대신, 거대한 '지식 도서관'을 따로 짓고, 필요한 때만 '유능한 사서'가 필요한 책만 골라오게 하여, 더 똑똑하고 잊지 않는 AI 를 만들었습니다."

이 기술은 AI 가 더 많은 지식을 배우면서도, 그 지식을 오래 기억하고 효율적으로 사용할 수 있게 해주는 새로운 방향을 제시합니다. 마치 우리가 공부를 할 때 단순히 암기하는 것보다, 잘 정리된 참고서와 색인을 활용하는 것이 훨씬 효과적인 것과 같은 원리입니다.

1. 문제 정의 (Problem)

변환기 (Transformer) 아키텍처는 시퀀스 모델링에 강력한 기반을 제공하지만, 훈련 중에 획득한 지식을 명시적이고 주소 지정이 가능한 (addressable) 메모리 메커니즘으로 저장하고 조직화하는 구조적 결여를 가지고 있습니다.

현재의 한계: 기존 시스템은 KV 캐시 유지/압축, 외부 정보 검색 (RAG), 또는 이전 숨은 상태의 재사용과 같은 추론 시 (inference-time) 메커니즘에 의존합니다. 이는 훈련 중 획득한 사실적 지식을 모델 내부에 명시적으로 저장하고 확장하는 데 한계가 있습니다.
확장성 문제: 대규모 메모리 뱅크에 대한 밀집 (dense) 크로스 어텐션은 계산 비용이 매우 높아 확장하기 어렵습니다.
지속 학습의 문제: 추가 학습 (Continual Training) 또는 지시 미세 조정 (Instruction Fine-Tuning, IFT) 시, 모델은 기존 지식을 잊어버리는 '파괴적 망각 (Catastrophic Forgetting)' 현상을 겪기 쉽습니다.

2. 방법론 (Methodology)

저자들은 학습 가능한 희소 메모리 뱅크 (Learnable Sparse Memory Banks) 를 도입하고, 이를 효율적으로 확장하기 위해 Mixture of Chapters (MoC) 아키텍처를 제안합니다.

A. 학습된 메모리 뱅크 (Learned Memory Bank)

구조: 무작위로 초기화되고 엔드 - 투 - 엔드 (end-to-end) 로 훈련되는 잠재 토큰 (latent tokens) 의 집합으로 구성된 메모리 뱅크를 변환기 레이어에 통합합니다.
접근 방식: 토큰 스트림이 쿼리 (Query) 를 생성하고, 메모리 뱅크가 키 (Key) 와 값 (Value) 을 제공하여 크로스 어텐션 (Cross-Attention) 을 통해 지식을 검색합니다. 이는 텍스트 기반 검색이 아닌 컴팩트한 잠재 공간 (latent space) 에서 내부 검색을 가능하게 합니다.

B. Mixture of Chapters (MoC) - 확장성 해결

대규모 메모리 뱅크에 대한 밀집 어텐션의 비용 문제를 해결하기 위해, 전문가 혼합 (Mixture-of-Experts, MoE) 에서 영감을 받은 챕터 기반 라우팅을 도입합니다.

챕터 분할: 전체 메모리 뱅크를 여러 개의 '챕터 (Chapters)'로 분할합니다. (예: 총 262,208 개의 메모리 토큰을 4,097 개의 챕터로 분할).
라우터 (Router): 입력 시퀀스별로 경량 라우터를 사용하여 관련성 높은 소수의 챕터 (Top-k) 만 선택합니다.
계산 효율성: 모든 메모리 토큰에 어텐션하는 대신 선택된 챕터에만 어텐션함으로써 계산 비용을 $O(L \cdot N_m)$ 에서 $O(L \cdot k \cdot T)$ 로 줄여 확장성을 확보합니다.

3. 주요 기여 (Key Contributions)

변환기 내 학습된 연관 메모리 뱅크: 텍스트 검색이 아닌 컴팩트한 잠재 공간에서 내부 검색을 가능하게 하는 지속적이고 주소 지정 가능한 메모리 뱅크를 제안했습니다.
확장 가능한 희소 접근 (Mixture of Chapters): 메모리를 챕터로 분할하고 경량 라우터를 통해 입력별 소수 챕터를 선택하는 MoC 아키텍처를 개발하여, 대규모 메모리 용량을 계산 비용 증가 없이 확장했습니다.
고정 연산량 하의 확장 행동 분석: 계산 최적화 (compute-optimal) 스캘링 논의와 병행하여, 고정된 FLOPs(iso-FLOP) 설정에서 표준 변환기 대비 메모리 모델의 성능을 평가했습니다.
지속 학습 하의 지식 유지 증명: 추가 학습 (IFT) 과정에서 명시적 메모리가 사실적 지식의 유지 (Retention) 를 개선하고 망각을 줄인다는 것을 실험적으로 입증했습니다.

4. 실험 결과 (Results)

실험은 사전 학습 (Pretraining) 과 지시 미세 조정 (IFT) 단계에서 표준 변환기 (Vanilla) 및 iso-FLOP 베이스라인과 비교하여 수행되었습니다.

사전 학습 성능: 96 억 토큰 (9.6B tokens) 기준 사전 학습에서 MoC 모델은 계산량이 동일한 (iso-FLOP) 밀집 변환기 베이스라인보다 더 낮은 검증 손실 (Validation Loss) 을 기록했습니다.
지시 미세 조정 (IFT) 및 망각 방지:
- 지식 기반 태스크: ARC-Challenge 및 BoolQ 와 같은 지식 집약적 태스크에서, 표준 베이스라인은 IFT 후 성능이 크게 저하되었습니다 (ARC-C 는 6.69%p 하락, BoolQ 는 6.24%p 하락).
- MoC 모델의 안정성: 반면 MoC 모델은 BoolQ 에서 거의 변화가 없었으며 (+0.24%), ARC-C 에서도 하락폭이 베이스라인의 절반 미만 (2.68%p) 으로 매우 작았습니다. 이는 메모리 뱅크가 추가 학습 중 사실적 지식을 '앵커 (Anchor)' 역할을 하여 망각을 방지함을 시사합니다.
메모리 뱅크 고정 (Freezing) 효과: IFT 단계에서 메모리 뱅크를 고정 (학습률 0) 하거나 매우 낮은 학습률로 훈련하더라도 성능 차이가 거의 없었습니다. 이는 사전 학습된 메모리가 후속 학습 단계에서도 재사용 가능함을 의미합니다.
역설적 관찰: IFT 후 사전 학습 분포 (pretraining distribution) 에 대한 평가에서는 MoC 모델의 손실이 베이스라인보다 높았으나, 지식 벤치마크 성능은 유지되었습니다. 이는 백본 (Backbone) 이 지시 정렬 (Instruction Alignment) 에 적응하는 동안, 메모리가 사실적 지식을 안정적으로 유지하는 역할 분리 (Separation of Roles) 가 발생했음을 시사합니다.

5. 의의 및 결론 (Significance)

새로운 확장 축 (Scaling Axis): 모델 파라미터 수나 계산량 (FLOPs) 만을 늘리는 기존 스캘링 법칙에 더해, 학습된 명시적 메모리를 새로운 확장 축으로 제안했습니다.
지식 유지의 해법: 대규모 언어 모델 (LLM) 의 지속적인 학습과 미세 조정 과정에서 발생하는 '파괴적 망각' 문제를 해결할 수 있는 아키텍처적 접근을 제시했습니다.
효율성과 성능의 균형: MoE 에서 영감을 받은 라우팅 메커니즘을 메모리 선택에 적용함으로써, 대규모 메모리 용량을 유지하면서도 계산 비용을 통제 가능한 수준으로 유지했습니다.

이 연구는 변환기 모델이 내부적으로 지식을 조직화하고 저장하는 능력을 강화하여, 더 큰 지식 용량과 지속적인 학습 능력을 갖춘 차세대 언어 모델 개발의 가능성을 열었습니다.