Stacked from One: Multi-Scale Self-Injection for Context Window Extension

이 논문은 동일한 LLM 레이어를 기반으로 하단 모델이 긴 입력을 다중 그레인 표현으로 압축하고 상단 모델이 이를 복원하는 '자기 주입 (self-injection)' 아키텍처를 제안함으로써, 기존 모델의 제한된 컨텍스트 창을 효율적으로 확장하면서도 메모리 사용량과 추론 지연을 획기적으로 줄이는 방법을 제시합니다.

핵심

🧠 문제: "책 한 권을 읽으려면 책방 전체를 가져와야 할까?"

기존의 인공지능 (LLM) 은 마치 기억력이 좋은 사람과 같습니다. 하지만 이 사람은 한 번에 읽을 수 있는 책의 두께가 정해져 있습니다. (예: 8,000 자까지만 기억함).
만약 이 사람에게 128,000 자가 넘는 거대한 백과사전이나 긴 소설을 읽히고 "이 책에서 특정 인물이 한 말은 뭐였어?"라고 물으면, 그 사람은 기억이 안 나거나 (성능 저하), 머리가 터져버립니다 (메모리 부족).

기존 해결책들은 두 가지였습니다:

1. 더 많이 공부시키기: 128,000 자짜리 책으로 다시 처음부터 훈련시키는 것. (시간과 돈이 너무 많이 듭니다.)
2. 책을 잘라내기: 중요한 부분만 남기고 나머지는 버리거나, 책장을 넘기면서 기억하는 방식. (하지만 중요한 내용을 놓치거나 속도가 느려집니다.)

---

💡 해결책: SHAREDLLM 의 "스마트한 도서관 사서" 시스템

이 논문이 제안한 SHAREDLLM은 거대한 책을 읽는 방식을 완전히 바꿉니다. 마치 현명한 도서관 사서와 지식인이 팀을 이루는 것과 같습니다.

1. 두 명의 팀원 (하단 모델 & 상단 모델)

이 시스템은 같은 뇌 (모델) 에서 나온 두 명의 팀원으로 구성됩니다.

• 하단 모델 (압축기/사서): 이 사람은 거대한 책 (입력 데이터) 을 한 번에 다 읽을 수는 없지만, 책을 조각조각 잘라내어 핵심 내용만 요약해 줍니다.
• 상단 모델 (해석자/지식인): 이 사람은 요약된 내용과 현재 질문을 받아서 최종 답변을 만듭니다.

2. '나무' 구조의 메모장 (Context Tree)

가장 재미있는 부분은 하단 모델이 책을 요약하는 방식입니다. 단순히 줄인 게 아니라, 나무 (Tree) 모양으로 정리합니다.

• 나무의 뿌리: 책의 전체 내용 (대략적인 개요).
• 나뭇가지: 책의 주요 장 (Chapter).
• 잎사귀: 중요한 문장이나 세부 사항.

"질문 (Query) 에 따라 잎만 따다!"
사용자가 "이 책에서 '비밀번호'가 뭐였지?"라고 물으면, 하단 모델은 나무 전체를 다 읽지 않고, '비밀번호'와 관련된 잎사귀 (중요한 부분) 만 골라내서 상단 모델에게 전달합니다.

• 비유: 도서관 사서가 "이 책에서 '비밀번호' 관련 내용만 찾아줘"라고 하면, 사서가 책 전체를 읽지 않고 목차와 색인을 보고 해당 페이지만 뽑아내서 주는 것과 같습니다.

3. '자기 주입 (Self-Injection)': 같은 뇌의 협력

이 두 팀원은 **완전히 같은 뇌 (모델)**에서 태어났습니다. 그래서 서로 다른 언어를 쓰지 않아도 됩니다.

• 하단 모델이 요약한 내용을 상단 모델이 바로 이해할 수 있도록, 가장 아래쪽 (저수준) 레이어에서 정보를 바로 전달합니다.
• 비유: 두 사람이 서로 다른 언어를 배우고 통역을 할 필요 없이, 같은 언어로 속삭이듯 정보를 주고받는 것입니다. 그래서 매우 빠르고 효율적입니다.

---

🚀 왜 이 기술이 특별한가요?

1. 기억력 확장 (8K → 128K+):
   이 모델은 8,000 자짜리 책으로만 훈련받았지만, 128,000 자가 넘는 긴 글을 읽어도 잘 처리합니다. 마치 8,000 자짜리 공부를 한 사서가, 나무 구조를 이용해 128,000 자짜리 도서관을 완벽하게 관리하는 것과 같습니다.

2. 압도적인 속도 (2~3 배 빠름):
   기존 방식들은 긴 글을 읽을 때 컴퓨터 메모리 (RAM) 를 다 써서 멈추거나 (OOM), 속도가 매우 느렸습니다. 하지만 SHAREDLLM 은 불필요한 정보를 미리 잘라내고 필요한 것만 전달하므로, 기존 방식보다 2~3 배 더 빠르고 메모리도 훨씬 적게 사용합니다.

3. 비용 절감:
   거대한 데이터를 다시 훈련시킬 필요가 없습니다. 이미 있는 모델 (예: LLaMA) 을 가져와서 위와 같은 '나무 정리' 방식만 추가하면 되므로, 훈련 비용이 거의 들지 않습니다.

---

📝 한 줄 요약

"SHAREDLLM 은 거대한 책 (긴 문맥) 을 읽을 때, 책 전체를 통째로 기억하려 애쓰지 않고, '나무' 모양으로 내용을 정리해 핵심만 골라내는 '스마트한 사서' 시스템을 도입하여, 인공지능이 긴 글을 빠르고 정확하게 이해하게 해주는 혁신적인 기술입니다."

이 기술 덕분에 앞으로 인공지능이 긴 논문, 긴 소설, 긴 회의록 등을 처리할 때 훨씬 더 효율적이고 저렴해질 것으로 기대됩니다.

기술 요약

1. 문제 정의 (Problem)

현대 대규모 언어 모델 (LLM) 의 가장 큰 병목 현상은 **제한된 컨텍스트 윈도우 (Context Window)**입니다.

• 현재 한계: 입력 텍스트가 모델의 컨텍스트 한계를 초과할 경우, 성능이 급격히 저하되거나 환각 (Hallucination) 이 발생합니다.
• 기존 해결책의 문제점:
  • 연속적 사전 학습 (Continual Pre-training): 긴 컨텍스트 데이터로 모델을 다시 학습시키는 것은 데이터 수집 및 연산 비용이 매우 높습니다.
  • 위치 인코딩 확장 (Positional Interpolation 등): "짧게 학습하고 길게 테스트 (Train short, test long)"하는 방식은 효율성이 낮고, 128K 토큰과 같은 극단적인 길이를 처리할 때 메모리 부족 (OOM) 이나 추론 속도 저하를 초래합니다.
  • 스트리밍/인코더 - 디코더 방식: 메모리 효율은 개선되었으나, FlashAttention 과 같은 고성능 어텐션 구현과 호환되지 않거나, 인코더와 디코더 간의 정렬 (Alignment) 을 위해 추가적인 학습 단계가 필요하여 비용이 증가합니다.

2. 제안 방법론: SHAREDLLM (Methodology)

저자들은 SHAREDLLM이라는 새로운 프레임워크를 제안합니다. 이는 단일 LLM 의 레이어를 공유하여 사용하는 자기 주입 (Self-injection) 메커니즘과 다중 스케일 컨텍스트 압축을 핵심으로 합니다.

2.1. 기본 아키텍처

• 구조: 하나의 LLM 을 두 개의 모델 (하위 모델, 상위 모델) 로 구성합니다.
  • 하위 모델 (Compressor): 긴 입력 컨텍스트를 작은 청크 (Chunk) 로 나누어 다중 그레인 (Multi-grained) 표현으로 압축합니다.
  • 상위 모델 (Decoder): 현재 실행 중인 컨텍스트 (질문 등) 를 입력으로 받으며, 하위 모델에서 압축된 정보를 통합하여 다음 토큰을 생성합니다.
• 자기 주입 (Self-injection): 두 모델은 동일한 사전 학습된 LLM 의 레이어에서 초기화됩니다. 정보 전달은 **가장 낮은 레이어 (Lowest layers)**에서만 발생하며, 긴 순전파 (Forward pass) 나 중복된 크로스 어텐션 연산을 우회합니다. 이로 인해 추가적인 학습 비용 (Warm-up 등) 없이 기존 체크포인트에서 바로 파인튜닝이 가능합니다.

2.2. 컨텍스트 트리 (Context Tree) 및 동적 검색

효율적인 정보 압축과 검색을 위해 트리 구조를 도입했습니다.

• 트리 구조: 입력 텍스트 청크를 루트 노드에서 시작하여 재귀적으로 분할합니다.
  • 동적 분할: 쿼리 (Query) 에 따라 관련성이 높은 노드는 세분화 (Fine-grained) 하고, 관련성이 낮은 노드는 거칠게 압축 (Coarse-grained) 합니다.
  • 쿼리 인식 정책 ($\pi$):
    • 언어 모델링: Λ-모양 (Lambda-shape) 패턴을 모방하여 오른쪽 분기를 선택합니다.
    • 지시 따르기 (Instruction-following): 쿼리와 의미적 유사도가 높은 분기를 선택합니다.
• KV 상태 압축: 선택된 노드 (Preserved Node) 에 대해 Key-Value(KV) 상태를 추출하고, 균일한 다운샘플링을 통해 길이를 압축합니다. 상단 레일일수록 높은 압축 비율을 적용하여 거시적 정보를, 하단 레일일수록 미세한 정보를 보존하는 계층적 구조를 가집니다.

2.3. 위치 인식 크로스 어텐션

• 압축된 청크들의 순서를 유지하기 위해 청크 수준의 위치 인덱스를 부여하고 Rotary Positional Embedding (RoPE) 을 적용합니다.
• 이를 통해 상위 모델의 쿼리가 압축된 컨텍스트 블록 간의 상대적 거리를 올바르게 인식하도록 합니다.

3. 주요 기여 (Key Contributions)

1. SHAREDLLM 아키텍처: 공유된 KV 메커니즘을 통해 최소한의 튜닝 가능한 파라미터로 효율적인 컨텍스트 윈도우 확장을 가능하게 하는 계층적 구조 제안.
2. 컨텍스트 트리 (Context Tree): 비정형적인 긴 컨텍스트를 coarse-to-fine(거칠게부터 세밀하게) 방식으로 표현하고, 쿼리 인식 동적 알고리즘을 통해 필요한 정보만 효율적으로 추출하는 데이터 구조 개발.
3. 성능 및 효율성 입증: 8K 토큰으로만 학습되었음에도 128K 토큰 이상의 입력에 대해 뛰어난 일반화 (Extrapolation) 능력을 보이며, 기존 스트리밍 및 인코더 - 디코더 방식 대비 메모리 사용량과 추론 속도를 획기적으로 개선함.

4. 실험 결과 (Results)

4.1. 언어 모델링 (Language Modeling)

• 외삽 능력: 8K 토큰으로 학습된 모델이 128K 토큰 입력에서도 퍼플렉시티 (Perplexity) 폭발 없이 안정적으로 작동했습니다.
• 비교 우위: YaRN, CEPE 등 기존 방법론보다 낮은 퍼플렉시티를 기록했습니다. 특히 CEPE 는 인코더 - 디코더 정렬을 위한 추가 학습이 필요했으나, SHAREDLLM 은 오프 - 더 - 쉘 (Off-the-shelf) 체크포인트로 직접 파인튜닝 가능하여 학습 비용을 대폭 절감했습니다.

4.2. 장문 이해 벤치마크 (Long-context Understanding)

• LongBench 및 InfiniBench: 단일/다중 문서 QA, 요약, 코드 완성 등 다양한 태스크에서 SOTA(최고 수준) 인 베이스라인 (Activation Beacon, StreamingLLM 등) 을 능가하거나 동급의 성능을 보였습니다.
• 중간 잘라내기 (Middle-truncation) 없이도 우수한 성능: 많은 기존 모델이 성능 향상을 위해 입력의 중간 부분을 잘라내는 방식을 사용했으나, SHAREDLLM 은 전체 컨텍스트를 효율적으로 처리하여 더 높은 정확도를 달성했습니다.

4.3. 효율성 (Efficiency)

• 메모리 및 속도: 128K 토큰 입력 시 YaRN 은 메모리 부족 (OOM) 을 겪는 반면, SHAREDLLM 은 안정적으로 실행되었습니다.
• 추론 속도: 스트리밍 방식보다 2 배, 인코더 - 디코더 방식보다 3 배 빠른 추론 속도를 기록했습니다. 이는 불필요한 레이어 통과를 피하고 FlashAttention 과 호환되는 구조를 유지했기 때문입니다.

5. 의의 및 결론 (Significance)

SHAREDLLM 은 비용 효율성과 성능 사이의 최적 균형을 달성했습니다.

• 학습 비용 절감: 별도의 대규모 사전 학습이나 복잡한 정렬 과정 없이, 기존 짧은 컨텍스트 LLM 을 바로 장문 처리 모델로 변환할 수 있습니다.
• 실용성: 메모리 효율이 뛰어나고 추론 속도가 빠르므로, 실제 응용 환경에서 긴 문서 처리가 필요한 시나리오 (법률 문서 분석, 긴 대화 기록 요약 등) 에 적용하기 매우 유리합니다.
• 확장성: 이 패러다임은 멀티모달 LLM 등 다른 아키텍처로도 확장 가능하여, 궁극적으로 "무한한 컨텍스트 처리"를 위한 중요한 기반 기술로 평가됩니다.

이 연구는 LLM 의 컨텍스트 윈도우 확장 문제를 해결하기 위해 단순한 위치 인코딩 수정을 넘어, 구조적 압축과 효율적인 정보 검색을 결합한 새로운 접근 방식을 제시했다는 점에서 의의가 큽니다.