기술 요약: 언어 모델의 기억과 언어를 위한 기억 모델

문제 정의

본 논문은 대규모 언어 모델 (LLM) 의 "기억" 능력에 대한 이해의 근본적인 공백을 다루고 있습니다. 기계 학습 모델은 입력 정보를 은닉층 벡터 임베딩에 저장하지만, 인과적 언어 모델에서 이러한 임베딩의 정보적 내용은 제대로 규명되지 않았습니다. 저자들은 다음 토큰 예측을 위해 최적화된 현재 인과적 언어 모델은 본질적으로 고충실도 기억 형성에 부적합하다고 주장합니다. 그 이유는 다음 토큰 예측이 다대일 (many-to-one) 의 비가역 함수이기 때문이며, 모델은 전체 입력 시퀀스가 아닌 즉시 다음 토큰을 예측하는 데 충분한 정보만 유지하면 되기 때문입니다. 결과적으로 이러한 모델이 하류 작업 (검색 또는 확장된 컨텍스트 창과 같은) 을 위해 입력 시퀀스를 기억 임베딩으로 압축할 때, 일반 언어 작업에 필요한 "임의 정보 접근"을 지원하지 못합니다.

방법론

1. 정보 유지 측정

모델 임베딩에 얼마나 많은 입력 정보가 유지되는지 정량화하기 위해, 저자들은 역변환 기반 측정 프레임워크를 제안합니다:

• 아키텍처: 테스트 중인 모델인 고정된 인코더와 학습 가능한 디코더를 짝짓습니다. 디코더는 인코더의 최종 은닉층 임베딩으로부터 원래 입력 시퀀스를 재구성하려고 시도합니다.
• 지표:
  • 엔트로피 비율 ($H_r$): 재구성의 교차 엔트로피 손실과 이론적 경계 (균일 분포 대 완벽한 재구성) 를 비교하여 입력 정보의 유지 비율을 측정합니다.
  • 토큰 정확도: 디코더가 올바르게 식별한 입력 토큰의 비율입니다.
• 역변환 가능성: 저자들은 ~96 V100 시간의 계산 예산 내에서 입력을 재구성할 수 있는 역변환 함수를 실제로 찾을 수 있는 경우를 모델이 "좋은 기억"을 가진 것으로 정의합니다.

2. 기억 모델 아키텍처

본 논문은 정보 유지와 정보 추출을 분리하도록 설계된 병렬화 가능한 기억 모델 아키텍처를 소개합니다:

• 인코더 - 디코더 구조: 인코더는 입력 토큰의 청크를 고정 크기의 기억 임베딩으로 처리합니다. 디코더는 이러한 임베딩 (선택적으로 새로운 토큰과 함께) 을 처리하여 후속 토큰을 예측합니다.
• 병렬화: 순환 기억 모델 (예: Bulatov et al., 2022) 과 달리, 이 아키텍처는 인코더가 훈련 및 추론 모두 동안 모든 입력 청크를 병렬로 처리할 수 있게 합니다.
• 훈련 전략:
  • 결합 목적 함수: 다음 토큰 예측 손실과 "복제" 손실 (모델이 기억에서 시퀀스를 재현해야 함) 의 선형 결합으로 훈련합니다.
  • 커리큘럼 학습: 디코더가 먼저 "빈 복제" 작업을 통해 토큰 입력 없이 정보 풍부한 임베딩을 처리하도록 학습한 후, 표준 다음 토큰 예측으로 훈련하는 2 단계 접근법입니다.
  • 고정된 인코더: 먼저 고충실도 오토인코더 인코더를 훈련한 후 이를 고정하고, 디코더가 이러한 임베딩을 활용하도록 훈련합니다.

3. 실험 설정

• 데이터셋: FineWeb-edu, FineMath 및 다양한 표준 벤치마크 (LongBench, IFEval 등).
• 모델: 다양한 크기의 Mixers 와 Transformers (소형 커스텀 모델부터 Llama 3.1, Qwen, BERT Large 까지) 에서 테스트했습니다.
• 작업: 인과적 다음 토큰 예측, 입력 재생성 (오토인코딩), 복제 작업, 그리고 기억 임베딩만으로 토큰을 재구성하는 빈 복제 작업입니다.

주요 결과

1. 인과적 모델은 기억이 부족함

• 낮은 정보 유지: 인과적 언어 모델과 검색 모델은 매우 낮은 정보 유지를 보입니다. 모델 크기와 계산량이 증가하더라도 인과적 임베딩에서 입력을 재구성하는 엔트로피 비율 ($H_r$) 과 토큰 정확도는 여전히 낮습니다 (예: $n_{ctx}=512$인 경우 약 5% 토큰 정확도).
• 비가역성: 인과적 모델의 임베딩은 입력 시퀀스를 재구성하기에 부족하며, 이는 다음 토큰 예측 목표가 포괄적인 기억 형성을 장려하지 않음을 확인시켜 줍니다.

2. 오토인코더는 거의 완벽한 기억을 달성함

• 높은 충실도: 입력 재생성을 위해 특별히 훈련된 오토인코더는 거의 완벽한 기억 형성을 달성하여 토큰 정확도가 90% 를 초과하고 엔트로피 비율이 1.0 에 근접합니다.
• 데이터 의존성: 이러한 고충실도 기억은 "사소하지 않습니다" (즉, 단순히 토큰을 벡터 슬롯에 매핑하는 것이 아님); 분포 외 데이터로 일반화되고 균일한 무작위 노이즈에서는 실패하여 의미 있는 데이터 의존적 표현을 학습함을 증명합니다.

3. 결합 목적 함수와 커리큘럼 학습의 성공

• 결합 훈련: 인과적 및 복제 목적 함수의 조합으로 훈련된 모델은 인과적 예측 능력을 유지하면서 정보 풍부한 기억을 형성하고 디코딩하는 방법을 성공적으로 학습합니다.
• 커리큘럼 학습: "빈 복제" 커리큘럼 (토큰 입력 없이 기억을 디코딩하는 방법을 먼저 학습) 은 결정적입니다. 이를 통해 디코더는 고정된 고충실도 오토인코더 임베딩을 효과적으로 활용할 수 있습니다. 이 커리큘럼이 없으면 디코더는 종종 풍부한 기억 임베딩을 활용하지 못하고 토큰 기반 예측으로 돌아가는 경우가 많습니다.
• 성능: 결합 목적 함수 모델은 순수 인과적 모델에 비해 정보 관련 벤치마크 (예: LongBench, Multi-Document QA) 에서 상당한 개선을 보이지만, 일반 언어 이해 벤치마크 (예: HellaSwag) 에서는 혼합되거나 중립적인 결과를 보입니다.

4. 아키텍처 효율성

• 병렬화: 제안된 병렬 인코더 - 디코더 아키텍처는 추론 계산 복잡도를 $O(n^2)$에서 약 $O(n^{4/3})$ (완전 병렬화의 경우 인코더는 $O(n^{2/3})$) 로 줄여, 첫 번째 토큰 도달 시간 및 메모리 사용량에서 상당한 이점을 제공합니다.
• 인코더 - 디코더 불일치: 본 논문은 "Mixer" 인코더와 "Transformer" 디코더는 종종 호환되지 않는 학습된 매니폴드를 가지며, 짝을 이룰 때 효율성이 떨어지므로 아키텍처별 정렬이 필요하다고 지적합니다.

중요성과 주장

본 논문은 기억 모델이 역사적으로 전체 컨텍스트 Transformer 보다 성능이 저조했던 이유에 대한 이론적 및 경험적 근거를 제공합니다: 인과적 훈련은 본질적으로 비가역적이므로 고충실도 기억 형성에 부적합합니다.

주요 기여 사항은 다음과 같습니다:

1. 규명: 규모와 관계없이 표준 LLM 임베딩이 상대적으로 적은 입력 정보를 포함함을 입증합니다.
2. 패러다임 전환: 정확한 기억 형성은 다음 토큰 예측을 넘어선 목적 함수, 구체적으로 정보 유지를 명시적으로 요구하는 결합 목적 함수나 커리큘럼 학습이 필요함을 제안합니다.
3. 아키텍처: 계산 효율성과 임의 입력 정보 저장 및 검색 능력을 균형 있게 갖춘 병렬화 가능한 인코더 - 디코더 기억 모델을 소개합니다.
4. 검색에 대한 함의: 이 발견들은 현재 임베딩 기반 검색 시스템이 임베딩 자체의 낮은 정보 밀도로 인해 제한받으며, 이론적 한계에 비해 임의 정보 조회 작업에서 종종 실패하는 이유를 설명합니다.

저자들은 다음 토큰 예측만으로는 기억에 불충분하지만, 모델을 훈련하여 "정보 풍부한" 기억을 형성하고 디코딩하게 할 수 있으며, 이는 계산 효율성과 임의 정보 접근 능력을 모두 제공하는 새로운 종류의 기억 모델을 가능하게 한다고 결론지었습니다.