RAT+: Train Dense, Infer Sparse -- Recurrence Augmented Attention for Dilated Inference

이 논문은 사전 학습된 밀집 어텐션 모델을 재학습 없이도 추론 시 희소화 (dilated) 패턴으로 유연하게 전환하면서도 긴 범위의 연결성을 유지하고 정확도 저하를 최소화하는 'RAT+'라는 새로운 아키텍처를 제안합니다.

핵심

1. 문제: "긴 문서를 읽을 때의 딜레마"

지금까지의 인공지능 (LLM) 은 긴 글을 읽을 때 두 가지 큰 고민이 있었습니다.

• 모든 것을 다 읽는 경우 (Dense Attention): 도서관에서 책 한 권을 읽을 때, 첫 페이지부터 마지막 페이지까지 모든 글자를 하나하나 꼼꼼히 읽는 것입니다.
  • 장점: 내용을 완벽하게 이해합니다.
  • 단점: 책이 두꺼울수록 (문장이 길어질수록) 시간이 너무 오래 걸리고, 기억해둬야 할 메모 (KV Cache) 가 폭발적으로 늘어납니다.
• 일부만 읽는 경우 (Sparse/Dilated Attention): 책의 매 10 페이지, 20 페이지마다 한 번씩만 훑어보는 것입니다.
  • 장점: 속도가 매우 빠르고 메모리도 적게 씁니다.
  • 단점: 중요한 내용을 놓치기 쉽습니다. 특히 기존 모델들은 이렇게 '간격 두고 읽는' 방식을 갑자기 적용하면, 이해도가 급격히 떨어지는 치명적인 오류가 있었습니다. (예: "아까 그 페이지에서 무슨 일이 있었지?"를 잊어버림)

2. 해결책: RAT+ 의 "기억력 강화 훈련"

RAT+ 는 이 문제를 해결하기 위해 두 가지 혁신적인 전략을 사용합니다.

전략 1: "완벽한 훈련, 유연한 실전" (Train Dense, Infer Sparse)

기존 방식은 "간격 두고 읽는 법"을 처음부터 따로 가르쳤다면, RAT+ 는 **처음부터 모든 페이지를 꼼꼼히 읽는 법 (완벽한 훈련)**으로 학습시킵니다.

• 비유: 마치 마라톤 선수가 처음부터 42km를 달리는 법을 배우는 것과 같습니다.
• 효과: 훈련이 끝난 후, 실제 경기 (추론) 에서는 "10km 마다 한 번만 체크하는 법"을 적용해도, 이미 42km를 달릴 줄 아는 근육 (지식) 이 있기 때문에 실력이 떨어지지 않습니다.
• 핵심: 별도의 모델을 여러 개 만들 필요 없이, 하나의 모델로 상황에 따라 속도를 조절할 수 있습니다.

전략 2: "연결고리" (Recurrence) - 가장 중요한 비법

그런데 왜 기존 모델들은 '간격 두고 읽기'를 하면 망했을까요? 바로 연결고리가 끊겼기 때문입니다.

• 비유: 책장을 넘길 때마다 이전 페이지의 내용을 잊어버리고 새로운 페이지만 보는 것입니다.
• RAT+ 의 해결책: RAT+ 는 **'재귀 (Recurrence)'**라는 기술을 도입했습니다. 이는 **"이전 페이지의 핵심 내용을 요약해서 다음 페이지로 전달하는 메모"**를 만드는 것과 같습니다.
  • 비록 10 페이지 건너뛰고 읽더라도, "1 페이지에서 중요한 내용이 10 페이지로 전달되어 20 페이지로 이어진다"는 연결고리가 있기 때문에 내용을 놓치지 않습니다.
  • 이 연결고리는 **능동적으로 학습 (Active Recurrence Learning)**되도록 훈련되어, 모델이 스스로 "어떤 정보를 기억해야 할지"를 배우게 합니다.

3. RAT+ 의 놀라운 성과

이 논문은 RAT+ 가 다음과 같은 성과를 냈다고 말합니다.

1. 속도와 정확도의 두 마리 토끼:
   • D=16 (16 페이지마다 읽기): 원래 속도보다 16 배 빠르면서도, 정확도는 거의 떨어지지 않습니다. (약 1% 만 감소)
   • D=64 (64 페이지마다 읽기): 속도는 64 배 빨라지는데, 정확도도 놀랍게도 2~3% 만 떨어집니다.
2. 유연한 적용:
   • 짧은 글을 읽을 때는 모든 페이지를 꼼꼼히 보고, 긴 글을 읽을 때는 간격 두고 읽는 식으로 상황에 따라 자동으로 조절할 수 있습니다.
   • 심지어 '가장 중요한 부분만 골라 읽는 (Top-k)' 방식에서도 기존 모델보다 훨씬 잘합니다.
3. 확장성:
   • 모델의 크기를 키울수록 (2.6B 파라미터), 이 '간격 두고 읽기' 방식의 정확도 손실은 더 줄어들었습니다. 즉, 모델이 클수록 더 똑똑하게 간격 읽기를 잘한다는 뜻입니다.

4. 요약: RAT+ 가 가져오는 변화

기존의 인공지능은 **"긴 글을 읽으려면 무조건 무겁고 느려야 한다"**는 고정관념이 있었습니다. 하지만 RAT+ 는 다음과 같이 바꿉니다.

"우리는 처음부터 모든 것을 완벽하게 배웠기 때문에, 필요할 때는 가볍게 스킵하며 읽어도 실수하지 않습니다. 마치 기억력이 좋은 사람이 긴 책을 읽을 때, 중요한 부분만 빠르게 훑어봐도 전체 흐름을 완벽하게 이해하는 것과 같습니다."

이 기술이 상용화되면, 긴 문서 분석, 긴 대화 기록 처리, 실시간 번역 등에서 속도는 엄청나게 빨라지고 비용은 획기적으로 줄어드는 새로운 시대가 열릴 것입니다.

기술 요약

1. 문제 정의 (Problem)

현대 LLM 의 표준 어텐션 메커니즘 (Self-Attention) 은 시퀀스 길이에 대해 이차 (Quadratic) 복잡도를 가지므로, 긴 문맥 처리 시 계산 비용 (FLOPs) 과 메모리 (KV Cache) 가 급격히 증가합니다. 이를 해결하기 위해 **확장된 어텐션 (Dilated Attention)**과 같은 구조적 희소화 기법이 연구되어 왔습니다.

• 기존 접근법의 한계:
  • Scratch-from-training: 희소 아키텍처를 처음부터 학습시키는 방식은 각 설정 (dilation size 등) 마다 별도의 모델을 학습해야 하므로 비용이 많이 들고 유연성이 떨어집니다.
  • Inference-time Sparsification: 사전 학습된 밀집 모델을 추론 시 희소 패턴으로 변환하는 방식은 **국소 윈도우 (Local Window)**나 Top-k 블록 어텐션에서는 어느 정도 효과가 있으나, **확장된 어텐션 (Dilated Attention)**으로 변환할 경우 심각한 성능 저하가 발생합니다.
  • 핵심 원인: 밀집 어텐션 모델은 확장된 패턴으로 변환될 때, 끊어진 연결 (disconnected connections) 로 인해 **완전한 수용 영역 (Receptive Field)**을 구성하지 못해 성능이 급격히 떨어집니다.

2. 방법론 (Methodology)

저자들은 이러한 문제를 해결하기 위해 재귀 (Recurrence) 메커니즘을 어텐션에 통합한 RAT+ 아키텍처를 제안합니다. RAT+ 는 밀집 어텐션의 능력을 유지하면서 추론 시 다양한 희소 패턴으로 유연하게 전환할 수 있도록 설계되었습니다.

핵심 기술 요소

1. 전체 시퀀스 재귀 (Full-sequence Recurrence):

   • 기존 RAT 모델은 청크 (Chunk) 단위로 재귀를 적용했으나, RAT+ 는 **전체 시퀀스 길이 (L=T)**에 걸쳐 재귀를 적용합니다.
   • 이는 Dilated Attention 이 필요한 수용 영역을 명시적으로 구성하여, 희소 추론 시에도 정보 손실을 방지합니다.

2. 능동적 재귀 학습 (Active Recurrence Learning, ARL):

   • 단순히 재귀를 추가하는 것만으로는 모델이 밀집 어텐션에 의존하여 재귀 기능을 제대로 학습하지 않는 '게으른 학습 (Lazy Learning)' 문제가 발생합니다.
   • 이를 해결하기 위해 Joint Training 전략을 사용합니다. 한 배치 내에서 **밀집 설정 (D=1)**과 **희소 설정 (D=64 등)**을 동시에 학습시켜, 모델이 재귀의 필요성을 명확히 인식하고 학습하도록 유도합니다.
   • 특히, 재귀가 효과적으로 작동해야 하는 최소 길이 (예: 64 토큰) 를 강제하여 재귀 능력을 활성화합니다.

3. 중첩된 청크 (Overlapped Chunk) 및 적응 (Adaptation):

   • 추론 시 dilation size 를 변경하더라도 재귀의 출력 분포가 일관되도록 **고정된 길이 (L=64)**의 중첩된 청크를 사용합니다.
   • 사전 학습된 밀집 모델을 특정 희소 패턴에 맞추기 위해 1B 토큰 정도의 짧은 적응 (Resolution Adaptation) 단계만 거치면 되며, 별도의 재학습이 필요 없습니다.

3. 주요 기여 (Key Contributions)

1. 새로운 아키텍처 RAT+ 제안: 밀집 학습과 희소 추론을 동시에 지원하는 단일 모델을 설계하여, Dilated Attention 의 성능 저하 문제를 재귀 메커니즘으로 해결했습니다.
2. 유연한 추론 지원: 하나의 모델로 D=1(밀집) 부터 D=64(고도 희소) 까지의 다양한 확장 크기, 로컬 윈도우, 하이브리드 구성 (레이어/헤드 단위 혼합) 을 유연하게 지원합니다.
3. Top-k Block Attention 성능 향상: 기존 어텐션 모델보다 Top-k 블록 어텐션 (Needle-in-a-Haystack 등) 에서 재귀가 블록 선택 알고리즘의 정확도를 높여 더 우수한 성능을 보입니다.
4. 확장성 입증: 1.5B 파라미터 (100B 토큰 학습) 및 2.6B 파라미터 (200B 토큰 학습) 모델에서 동일한 유효성을 입증했으며, 모델 크기가 커질수록 밀집과 희소 간의 성능 격차가 줄어듦을 확인했습니다.

4. 실험 결과 (Results)

• 성능 유지: 1.5B 모델 기준, D=16 설정에서 밀집 어텐션과 거의 동일한 정확도를 유지하며, D=64 설정에서도 상식 추론 (Commonsense Reasoning) 과 LongBench 태스크에서 약 2~3 점의 미세한 성능 하락만 발생했습니다.
• 효율성 극대화:
  • FLOPs 및 KV Cache 감소: D=16 설정 시 어텐션 계산량과 KV Cache 크기를 16 배 감소시킵니다.
  • 처리량 (Throughput) 향상: 4K 컨텍스트 길이에서 D=16 시 최대 10 배, 16K 컨텍스트에서는 20 배 이상의 처리량 향상을 달성했습니다.
  • LongBench: 긴 문맥 이해 태스크에서 Dilated Attention 과 Local Window 를 결합한 하이브리드 설정이 우수한 성능을 보였습니다.
• Needle-in-a-Haystack (NIAH): RAT+ 는 희소화된 Top-k 블록 어텐션에서도 밀집 모델 대비 훨씬 높은 정확도를 유지하며, 재귀가 블록 선택의 정확도를 높이는 데 기여함을 입증했습니다.

5. 의의 및 결론 (Significance)

RAT+ 는 효율적인 언어 모델 개발에 중요한 방향성을 제시합니다.

• 유연성과 효율성의 조화: 별도의 희소 모델을 학습할 필요 없이, 단일 밀집 모델을 학습한 후 추론 시 상황에 따라 최적의 희소 패턴을 선택할 수 있게 합니다. 이는 하드웨어 제약이나 태스크 요구사항에 따라 동적으로 리소스를 조절할 수 있는 실용적인 솔루션입니다.
• 재귀의 부활: Transformer 의 효율성 문제를 해결하기 위해 RNN 의 재귀 메커니즘을 현대적인 어텐션 구조에 성공적으로 통합하여, 긴 시퀀스 모델링에서 재귀가 가진 수용 영역 구성 능력을 활용했습니다.
• 미래 지향성: 이 연구는 토큰화 없는 (Byte-level) 입력 처리, 추론 단계에서의 샘플링 전략 (탐색은 희소, 생성은 밀집) 등 다양한 응용 분야로 확장될 가능성을 열어줍니다.

요약하자면, RAT+ 는 **"한 번 학습한 모델로 다양한 효율성 요구사항을 충족한다"**는 개념을 실현하여, 대규모 언어 모델의 배포 비용과 추론 속도를 획기적으로 개선할 수 있는 강력한 아키텍처입니다.