본문 바로가기

논문 리뷰/Language Model

LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens

[arXiv](2024/02/21 version v1)

 

 

Abstract

검색을 통해 최적의 RoPE 재조정 인자를 찾고 LLM context window를 2048k까지 확장

 

 

 

Non-uniformity in Positional Interpolation

Preliminary

Rotary Position Embedding

임베딩 차원 d, 토큰 위치 n, 회전 주파수(=회전율, 회전 각도. 차원 i가 작을수록 주파수와 회전각이 크다)가 다음과 같을 때,

RoPE는 다음과 같이 단순화될 수 있다.

 

Context window extension ratio s and positional interpolation

원래 context 길이 L에 대한 확장된 context 길이 L'의 비율 s를 정의한다. ( s = L'/L )

 

Context를 확장하는 방법으로는 외삽과 보간이 있다.

  • Extrapolation: 훈련 시의 최대 context 길이를 초과하는 인덱스를 사용한다.
  • Interpolation: 예를 들어, 최대 context 길이가 20인데 40개의 토큰을 입력하고 싶다면, 정수 인덱스를 서로 보간하여 0.5 인덱스를 추가해 사용하는 것이다. e.g. 1, 1.5, ..., 19.5, 20

 

RoPE에서 context를 확장하기 위한 위치 보간 방법은 회전 주파수를 축소하는 것이다.

s와 관련된 재조정 인자 λ,

일 때,

 

 

Linear positional interpolation (PI)

λ = s. 모든 회전율이 선형적으로 감소한다. 

위치 정보가 매우 밀집되어 인접한 토큰을 구별하기 힘들기 때문에 높은 s에서 성능이 저하된다.

 

NTK-based interpolation and extrapolation

λ = si. 낮은 차원(고주파수)에서는 적게, 높은 차원(저주파수)에서는 많이 스케일링하여 보간과 외삽이 모두 발생한다.

Fine-tuning 없이도 적용할 수 있다는 장점. 

고주파수에서 더 많이 스케일링해야 되는 것 아닌가? 라고 생각했지만 고주파수에서는 미세한 각도 차이를 구별하지 못하니 잘 구별하는 저주파수에서 더 많이 스케일링해야 되는 것인가? 

 

YaRN

고주파수에서는 외삽 (λ = 1), 저주파수에서는 PI, 중간 주파수에서는 NTK를 사용한다.

핵심은 경험적으로 RoPE 차원을 그룹화하는 것에 있으며 이로 인해 최적인 결과가 나오지 않을 수 있다.

 

Study on Non-uniform Positional Interpolation

실증적 분석을 통해 얻은 주요 결과:

  1. RoPE 차원은 상당한 불균일성을 나타내며 현재까지의 보간 방법으로는 효과적으로 처리할 수 없다.
  2. StreamingLLM에서 주목한 바와 같이 attention sink 현상으로 인해 초기 토큰에 대한 보간은 좋지 않다.
  3. 불균일성을 잘 활용하여 보간하면 효과적으로 context를 확장할 수 있다.

 

 

 

LongRoPE

Problem Formulation

우리의 목표는 perplexity를 최소화하는 재조정 인자 집합 I(λ̂, n̂)을 찾는 것이다.

n̂은 임계값으로 각 차원의 초기 n̂개의 인덱스에는 재조정을 적용하지 않는다.

 

Searching the Non-uniform Position Interpolation

Search space

 

Evolution-based search

s가 커질수록 기하급수적으로 검색 공간이 증가한다. 

Evolution search를 사용하며 효과적인 검색을 위한 2가지 추가 최적화 기술을 소개한다.

 

1. 최적화된 초기 인구 생성

P개의 재조정 인자를 무작위로 초기화하는 대신 PI, NTK, YaRN에 해당하는 재조정 인자를 초기 인구에 추가하고 나머지 P-3 개체는 3가지의 각 재조정 인자를 무작위로 변이 시킨다.

 

2. 단조 증가 제약

초기 인구에 대해 입력 X에 대한 perplexity를 계산하고 상위 k개의 개체를 선택한다. 

 

선택된 개체의 자손을 번식시킬 때 검색 공간을 줄이기 위하여 모든 차원에 대해 재조정 인자가 단조 증가하는, 즉 λi ≤ λi+1인 개체만을 생성한다.

 

8× extension without fine-tuning

Fine-tuning을 하지 않고도 성능 저하 없이 context 길이를 8배까지 확장할 수 있다.

 

Extending LLM Context Window to 2048K

Progressive extension to 2048k

LongRoPE 검색을 통해 128k로 context 확장 → 128k LLM fine-tuning → 256k context 확장 → 256k fine-tuning → 2048k context 확장

 

훈련 비용과 안정성 문제로 2048k에 대한 LLM fine-tuning은 수행하지 않는다. 

 

Shorter context window recovery

Context를 확장하면 원래의 짧은 context에 대한 성능이 저하되는 문제가 있다.

 

추가적으로 짧은 context에 대한 evolution search를 수행하고 추론 중에 이를 활용해 재조정 인자를 동적으로 조정한다.

 

 

 

Experiments