RePo: Language Models with Context Re-Positioning

이 논문은 인지 부하 이론을 기반으로 컨텍스트의 고정된 순서 구조를 완화하여 불필요한 인지 부하를 줄이고, OLMo-2 모델을 기반으로 한 RePo 메커니즘을 통해 노이즈가 많은 맥락이나 긴 문맥에서도 성능을 향상시키는 새로운 언어 모델 접근법을 제안합니다.

핵심

🧠 핵심 아이디어: "책상 정리하기"

인공지능이 글을 읽거나 문제를 풀 때, 마치 책상 위에 모든 자료를 펼쳐놓고 작업하는 사람과 같습니다.

1. 기존 방식 (기존 LLM):

   • 책상 위에 들어온 자료 (문장) 를 들어온 순서대로 일렬로 늘어놓습니다.
   • "첫 번째 문장, 두 번째 문장, 세 번째 문장..." 식으로 번호를 매겨서 처리합니다.
   • 문제점: 만약 책상 위에 **불필요한 잡동사니 (노이즈)**가 너무 많거나, **중요한 정보가 책상 끝쪽 (긴 문서)**에 숨겨져 있다면?
     • 인공지능은 "순서대로"만 보느라 머리가 피로해집니다. (인지 부하 증가)
     • 중요한 정보를 찾기 위해 책상 전체를 뒤져야 하므로, 진짜 중요한 생각 (추론) 을 할 에너지가 부족해집니다.

2. 새로운 방식 (REPO):

   • REPO 는 인공지능에게 **"순서대로 놓지 말고, 중요도에 따라 책상을 재배치해라!"**라고 가르칩니다.
   • 들어온 정보를 보고, **"이건 핵심이야! (가까이 놓자)", "이건 그냥 잡동사니야! (멀리 치자)"**라고 스스로 판단하여 위치를 다시 정합니다.
   • 마치 명품 가게의 진열사처럼, 가장 중요한 보석 (핵심 정보) 을 가장 잘 보이는 곳에, 나머지는 구석에 배치하는 것과 같습니다.

---

🚀 REPO 가 해결하는 3 가지 문제

이 기술은 특히 다음과 같은 상황에서 인공지능의 실력을 비약적으로 높여줍니다.

1. "Haystack(건초더미) 속의 바늘 찾기" (Noisy Context)

• 상황: 수천 줄의 긴 문서 속에 정답이 딱 하나 숨겨져 있고, 나머지는 다 헛소리일 때.
• 기존 AI: "순서대로" 읽다가 지쳐서 정답을 놓칩니다.
• REPO: "아, 이 부분이 정답이겠구나!" 하고 정답이 있는 부분을 가까이 끌어당겨 집중합니다. 잡동사니는 무시하고 정답만 쏙 집어냅니다.

2. "표 (Table) 데이터" (Structured Data)

• 상황: 복잡한 표나 차트를 글자로 바꿨을 때 구조가 깨지는 경우.
• 기존 AI: 글자 순서만 따라가서 "행과 열"의 관계를 잊어버립니다.
• REPO: 표의 구조를 이해하고, 관련된 정보들을 그룹화해서 배치합니다. 마치 표를 다시 그려서 보기 쉽게 만드는 것과 같습니다.

3. "긴 문서 읽기" (Long Context)

• 상황: 책 한 권 분량의 글을 읽고 마지막에 질문을 받으면.
• 기존 AI: 처음에 읽은 내용은 기억이 흐려집니다. (기억력 한계)
• REPO: 긴 문서에서도 중요한 부분과 질문 사이의 거리를 줄여줍니다. 마치 긴 줄을 접어서 중요한 부분끼리 붙여놓는 것처럼, 기억하기 쉽게 만들어줍니다.

---

🧩 어떻게 작동할까요? (비유: 지능형 도서관 사서)

기존 AI 는 책이 들어오는 **순서 (1 번, 2 번, 3 번)**만 따지는 자동 기계였습니다.

하지만 REPO는 지능형 도서관 사서처럼 작동합니다.

• 독자가 "이 책의 300 페이지에 있는 내용이 중요해!"라고 말하면, 사서는 300 페이지를 바로 찾아서 독자의 책상 바로 앞에 가져다 줍니다.
• 그리고 나머지 1~299 페이지는 뒤로 밀어둡니다.
• 이렇게 하면 독자는 **책상 (작업 기억력)**이 넓어지고, 생각할 에너지를 더 많이 쓸 수 있게 됩니다.

이 기술은 인공지능이 자신의 뇌 (작업 기억력) 를 아껴서, 더 깊은 생각 (추론) 을 할 수 있게 해줍니다.

---

📊 결과는 어떨까요?

연구진은 이 기술을 적용한 인공지능을 테스트해 보았습니다.

• 잡동사니가 많은 긴 글을 읽을 때 정답률이 훨씬 높아졌습니다.
• 표나 차트를 이해하는 능력도 크게 향상되었습니다.
• 짧은 대화를 할 때는 기존 방식과 똑같이 잘하지만, 긴 문서가 필요할 때는 압도적으로 잘합니다.

💡 결론

REPO는 인공지능에게 "무조건 순서대로 읽지 말고, 무엇이 중요한지 스스로 판단해서 정보를 재배치하라"는 새로운 능력을 선물한 기술입니다.

이제 인공지능은 더 이상 책상 위에 쌓인 서류 더미에 압도되지 않고, 중요한 정보를 쏙쏙 골라내어 더 똑똑하고 정확한 답변을 줄 수 있게 되었습니다. 앞으로 긴 문서 요약, 복잡한 데이터 분석, 그리고 더 긴 대화를 하는 AI 들의 핵심 기술이 될 것으로 기대됩니다.

기술 요약

1. 문제 정의 (Problem)

현대 대규모 언어 모델 (LLM) 의 핵심 기능인 인-컨텍스트 학습 (In-context Learning) 은 제한된 컨텍스트 윈도우 내에서 정보를 처리하는 인간 작업 기억 (Working Memory) 과 유사합니다. 그러나 기존 LLM 아키텍처는 다음과 같은 한계를 가지고 있습니다.

• 경직된 위치 인코딩: 대부분의 모델은 토큰에 0 부터 $L-1$ 까지의 선형 (linear) 이나 일정한 (constant) 정수 인덱스를 고정적으로 할당합니다.
• 인지 부하 (Cognitive Load) 증가: 인지 부하 이론 (Cognitive Load Theory, CLT) 에 따르면, 정보가 조직화되고 제시되는 방식에서 발생하는 불필요한 비용은 '외재적 부하 (Extraneous Load)'를 증가시킵니다.
• 작업 기억의 낭비: 이러한 불필요한 외재적 부하는 모델의 유한한 작업 기억 용량을 소모하여, 실제 추론과 주의 집중 (Attention Allocation) 에 필요한 '유용한 부하 (Germane Load)'를 처리하는 능력을 저하시킵니다.
• 성능 저하: 이로 인해 긴 컨텍스트, 노이즈가 많은 데이터, 또는 구조화된 데이터 (예: 표) 와 같은 복잡한 작업에서 성능이 크게 떨어지는 현상이 발생합니다.

2. 제안 방법: REPO (Methodology)

저자들은 REPO (Context Re-Positioning) 라는 새로운 메커니즘을 제안하여 토큰의 위치를 고정된 선형 순서가 아닌, 토큰 간의 관련성 (Relevance) 에 기반하여 동적으로 재배치합니다.

• 핵심 모듈 ($f_\phi$):
  • 각 토큰의 은닉 상태 (Hidden State) 를 입력받아 연속적이고 비선형적인 공간에서 새로운 위치 값 ($z_i$) 을 할당하는 경량 가분산 (Differentiable) 모듈입니다.
  • 위치 표현 (Position Representation): 토큰의 은닉 상태에서 위치 정보를 추출하기 위해 경량 SwiGLU 서브레이어를 사용합니다.
  • 위치 할당 (Position Assignment): 추출된 표현을 선형 변환을 통해 실수 값의 위치 $z_i$ 로 매핑합니다.
• 작동 원리:
  • 기존 RoPE(Rotary Position Embedding) 와 같은 가분산 위치 인코딩 함수 $g_\theta$ 를 사용합니다.
  • 주의 점수 (Attention Score) 계산 시, 토큰 간의 거리 $j-i$ 대신 예측된 위치 값의 차이 $z_j - z_i$ 를 사용합니다.
    $$A^{REPO}_{i,j} = q_i^\top g_\theta(z_j - z_i) k_j$$
  • 이 모듈은 LLM 의 특정 레이어 (예: 1/3 지점 이후) 부터 적용되며, 각 어텐션 헤드마다 독립적으로 학습됩니다.
• 효율성:
  • 토큰의 자동 회귀적 생성 순서 (Auto-regressive order) 는 변경하지 않고, 오직 어텐션 계산 시의 위치 인코딩에만 영향을 줍니다. 따라서 KV 캐시 재계산 없이 기존 모델에 적용 가능하며, 파라미터 증가량은 0.9% 수준으로 매우 경량입니다.

3. 주요 기여 (Key Contributions)

1. 인지 부하 이론 기반의 새로운 접근: LLM 의 위치 인코딩을 단순한 인덱스 할당이 아닌, 컨텍스트 구조를 이해하고 재구성하는 학습 가능한 과정으로 재정의했습니다.
2. 유연한 위치 패턴 학습: REPO 는 고정된 선형 패턴이나 상수 패턴에 국한되지 않고, 컨텍스트에 따라 상수 (Constant), 단조 (Monotonic), 또는 하이브리드 (Hybrid) 패턴을 동적으로 학습합니다.
3. 광범위한 실험 검증: 오픈 소스 모델인 OLMo-2 (1B 및 7B) 를 기반으로 계속적인 사전 학습 (Continual Pre-training) 을 수행하여 데이터 오염 문제를 배제하고 엄격하게 검증했습니다.

4. 실험 결과 (Results)

REPO 는 노이즈가 많은 컨텍스트, 구조화된 데이터, 긴 컨텍스트 작업에서 기존 베이스라인 (RoPE, NoPE 등) 을 압도적으로 능가했습니다.

• 노이즈가 많은 컨텍스트 (Noisy Context):
  • RULER 벤치마크 (Needle-in-a-Haystack): OLMo-2 1B 모델에서 REPO 는 평균 정확도가 **91.3%**로, 기존 RoPE 대비 5.4 포인트 향상되었습니다. 특히 중요한 'Needle' 토큰에 더 많은 주의를 기울이는 것을 확인했습니다.
• 구조화된 데이터 (Structured Data):
  • HybridQA (표 데이터): 표를 자연어로 변환할 때 발생하는 구조 정보 손실을 보완하여, OLMo-2 1B 에서 2.27 포인트, 7B 에서 4.09 포인트의 정확도 향상을 보였습니다.
• 긴 컨텍스트 (Longer Context):
  • LongBench 및 RULER (4K~16K 토큰): 훈련 시 보지 못한 긴 컨텍스트 (8K, 16K) 에서도 REPO 는 RoPE 보다 월등히 좋은 일반화 성능을 보였습니다. OLMo-2 1B 기준 LongBench 평균 점수가 6.93 포인트 향상되었습니다.
• 일반 작업 (General Tasks):
  • ARC, MMLU-Pro 등 일반적인 단거리 질문 답변 벤치마크에서는 기존 RoPE 와 유사한 성능을 유지하며, REPO 가 일반 성능을 해치지 않음을 입증했습니다.

5. 분석 및 통찰 (Analysis)

• 주의 분포 (Attention Mass): REPO 는 먼 거리에 있지만 중요한 'Needle' 토큰에 더 많은 주의를 집중시키고, 가까운 'Query' 토큰에 대한 과도한 주의 (Locality Bias) 를 줄이는 것을 확인했습니다.
• 학습된 위치 패턴:
  • 할당된 위치 값은 밀집된 비선형 공간에 분포하며, 컨텍스트의 내재적 구조 (예: Few-shot 예제의 구분) 를 포착합니다.
  • 모델은 컨텍스트에 따라 NoPE(상수 위치) 와 RoPE(선형 위치) 의 장점을 혼합한 하이브리드 패턴을 학습합니다.
• 효율성: 추론 시간과 FLOPs 증가가 미미하여 (약 0.9% 파라미터 증가) 실제 적용에 매우 유리합니다.

6. 의의 및 결론 (Significance)

이 논문은 LLM 이 컨텍스트 정보를 처리하는 방식을 수동적이고 경직된 구조에서 능동적이고 적응적인 구조로 전환할 수 있음을 보여줍니다.

• 아키텍처 혁신: 위치 인코딩을 사전 정의된 규칙이 아닌, 모델이 학습하여 컨텍스트의 구조와 의존성을 최적화하는 과정으로 바꿈으로써, 작업 기억 용량을 더 효율적으로 활용하게 합니다.
• 실용적 가치: 긴 문서 이해, 검색 증강 생성 (RAG), 에이전트 시스템 등 복잡한 컨텍스트가 필요한 실제 응용 분야에서 모델의 신뢰성과 성능을 크게 향상시킬 수 있는 잠재력을 가집니다.
• 개방성: 코드와 모델 가중치를 오픈소스하여 (SakanaAI), 연구 커뮤니티가 이를 확장하고 검증할 수 있는 기반을 마련했습니다.

요약하자면, REPO는 인지 과학의 통찰을 LLM 아키텍처에 적용하여, 모델이 컨텍스트의 '의미'에 따라 토큰의 위치를 재배치함으로써 복잡한 추론 능력을 획기적으로 개선한 획기적인 연구입니다.