🤖 machine learning

Jet-Long: Efficient Long-Context Extension with Dynamic Bifocal RoPE

Jet-Long은 짧은 문맥 충실도와 긴 문맥 외삽 사이의 균형을 맞추기 위해 RoPE 리스케일링 인자를 동적으로 조정하는 튜닝이 필요 없는 제로샷 방법으로, 효율적인 바이포컬 어텐션 메커니즘을 통해 최소한의 추론 오버헤드로 긴 문맥 벤치마크에서 최첨단 성능을 달성합니다.

원저자: Haozhan Tang, Zerui Wang, Yuxian Gu, Song Han, Han Cai

게시일 2026-07-10

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Haozhan Tang, Zerui Wang, Yuxian Gu, Song Han, Han Cai

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신에게 어떤 특정 길이(예를 들어 32,000 단어)까지의 이야기를 읽도록 훈련된 아주 똑똑한 로봇 뇌(거대 언어 모델)가 있다고 상상해 보세요. 그런데 당신이 이 로봇에게 그보다 4배나 더 긴 소설 한 권이나 거대한 코드 저장소를 읽게 하고 싶어 합니다. 만로 단순히 긴 텍스트를 건네주면 로봇은 혼란에 빠집니다. 이는 마치 책의 페이지 번호가 100에서 갑자기 1,000,000으로 점프하는 것을 보는 것과 같습니다. 로봇의 내부 나침반(RoPE라고 불리는 것)이 격렬하게 회전하며, 로봇은 환각을 일으키거나 이야기의 중간 부분을 잊어버리기 시작합니다.

한동안 과학자들은 로봇의 시야를 확장하기 위해 하나의 "마법의 숫자"를 선택하여 이를 해결하려고 노력했습니다. 하지만 이것은 승자 없는 게임이었습니다. 만약 끝까지 보기 위해 시야를 너무 많이 늘리면 로봇은 앞부분을 잊어버렸습니다. 만약 앞부분을 기억하기 위해 시야를 좁게 유지하면 끝까지 볼 수 없었습니다.

여기에 등장한 Jet-Long은 로봇의 눈에 동적 이중 초점 렌즈(dynamic bifocal lens) 역할을 하는 새로운 방법입니다.

이중 초점 안경의 기술

로봇에게 책 전체에 대해 하나의 확장된 시야를 강요하는 대신, Jet-Long은 두 개의 안경을 동시에 제공합니다:

근거리 렌즈: 텍스트의 첫 번째 덩어리(로컬 윈도우)를 위해, 로봇은 훈련받은 그대로의 모든 것을 봅니다. 확장이나 왜곡이 없습니다. 이를 통해 로봇은 최근의 과거에 대한 기억을 날카롭게 유지합니다.
원거리 렌즈: 나머지 텍스트(리모트 윈도우)를 위해, 로봇은 특별한 기술을 사용합니다. 단순히 시야를 늘리는 것이 아니라, 동적으로 페이지들을 그룹화합니다. 책이 길어짐에 따라, 로봇은 얼마나 많은 페이지를 하나의 "슈퍼 페이지"로 묶을지를 자동으로 조정합니다.

마법은 이 그룹화 계수가 고정되어 있지 않다는 점에 있습니다. 이는 현재 텍스트의 길이에 따라 실시간으로 변합니다. 텍스트가 짧으면 로봇은 모든 단어를 개별적으로 봅니다. 텍스트가 매우 길어지면, 로봇은 내부 나침반이 차분하게 유지되고 통제력을 잃지 않도록 적절히 단어들을 그룹화합니다. 즉, Jet-Long은 짧은 작업에서는 완벽한 정확도를 유지하면서도, 방대한 문서를 길을 잃지 않고 읽을 수 있게 해줍니다.

속도의 "공짜 점심"

보통 두 가지 다른 계산(가까이 보기와 멀리 보기)을 동시에 수행하면 로봇이 두 배로 느려질 것입니다. 하지만 저자들은 이 두 계산을 결합하는 영리한 방법을 찾아냈습니다. 그들은 두 렌즈의 수학적 계산을 한 번의 과정으로 처리하는 특별한 엔진("퓨즈드 커널", fused kernel)을 구축했습니다.

이는 마치 보통 채소를 썰고, 삶고, 볶는 세 가지 과정을 각각 별도의 냄비에서 수행해야 하는 요리사가, 세 단계를 한 번에 수행하면서도 맛을 전혀 잃지 않는 마법의 냄비를 가진 것과 같습니다.

결과: 로봇이 128,000 토큰의 방대한 컨텍스트를 읽을 때, Jet-Long은 H100 칩에서 표준 방식(FlashAttention-2)보다 시작 단계(프리필 단계)에서 1.39배 더 빠릅니다.
주의점: 로봇이 한 번에 한 단어씩 새로운 텍스트를 생성할 때는 표준 방식보다 약 4% 정도 느립니다. 이는 도서관 전체를 읽을 수 있게 되는 것에 비하면 아주 작은 대가입니다.

테스트 내용 (그리고 하지 않은 것)

팀은 세 가지 크기의 로봇 뇌(17억, 40억, 80억 파라미터)를 대상으로 테스트를 진행했으며, Jet-Long이 이전의 최고 방법들을 일관되게 이겼다는 것을 발견했습니다.

로봇이 텍스트의 건초더미 속에서 숨겨진 바늘을 찾을 수 있는지 확인하는 RULER 테스트에서, Jet-Long은 가장 작은 모델에서 기존 최고 경쟁자보다 4.79점 높게, 가장 큰 모델에서는 2.03점 높게 점수를 기록했습니다.
실제 검색 작업을 시뮬레이션하는 HELMET-RAG 테스트에서, Jet-Long은 전반적으로 최고의 정확도를 달 achievement 했습니다.
오래된 책을 읽는 PG-19 테스트에서, 다른 방법들이 텍ema가 길어짐에 따라 로봇의 혼란(퍼플렉서티)을 급증하게 만든 반면, Jet-Long은 혼란을 가장 낮은 수준으로 유지했습니다.

중요 참고 사항: 논문은 로봇을 작동시키기 위해 처음부터 다시 훈련시킬 필요가 없다는 아이디어를 명시적으로 배제합니다. Jet-Long은 "제로샷(zero-shot)" 모델에서 작동합니다. 즉, 이미 훈련된 로봇을 가져와서 이 렌즈를 그냥 끼워 넣기만 하면 됩니다. 그러나 논문은 Jet-Long이 "회전하는 나침반" 문제는 해결하지만, "어텐션 확산(attention diffusion, 로봇이 너무 많은 옵션에 압도되는 현상)"이라는 다른 문제는 해결하지 못한다고 언급했습니다. 이 문제는 렌즈를 추가하는 것이 아니라 로봇의 구조 자체를 완전히 바꾸어야 해결되는 문제입니다.

결론

Jet-Long은 로봇을 무한히 똑똑하게 만드는 마법 지팡이는 아니지만, 기존의 로봇들이 정신을 잃지 않고 최대 128,000 토큰 길이의 텍스트를 처리할 수 있게 해주는 매우 효율적이고 튜닝이 필요 없는 도구입니다. 이는 단기 기억을 완벽하게 유지하면서도 장기적인 시야를 동적으로 조정하며, 표준 방식만큼 빠르게 실행됩니다. 저자들은 실제 하드웨어(H100 GPU)에서 이를 측정했으며, 다양한 모델 크기와 서로 다른 유형의 어텐션이 섞인 하이브리드 로봇 뇌에서도 작동함을 확인했습니다. 이는 더 긴 글을 읽게 만들려는 모든 이들에게 검증된 강력한 업그레이드입니다.

기술 요약: Jet-Long

문제 정의

현대적인 대규모 언어 모델(LLM)은 검색 증강 생성(RAG), 저장소 수준의 코딩, 에이전트 워크플로우와 같은 긴 문맥(long-context) 애플리케이션에 점점 더 많이 배치되고 있습니다. 이러한 시나리오에서 누적된 추론 흔적(reasoning traces)과 도구 출력값은 종종 입력 길이를 사전 학습 윈도우(예: 32K 학습 윈도우 대비 100K+ 토큰)보다 한 자릿수 이상 초과하게 만듭니다.

긴 문맥을 직접 학습시키는 것은 이차적인 연산 비용과 부족한 긴 문맥 데이터로 인해 매우 비용이 많이 듭니다. 또한 파인튜닝은 짧은 문맥 성능을 저하시키는 경우가 많습니다. 결과적으로, 제로샷 문맥 확장(파인튜닝 없이 문맥을 확장하는 것)이 오픈 웨이트 체크포인트의 지배적인 배포 경로가 되었습니다.

그러나 회전 위치 임베딩(RoPE)을 사용하는 일반적인 트랜스포머 기반 LLM은 두 가지 주요 실패 모드로 인해 학습 윈도우를 넘어 일반화하는 데 실패합니다:

위치 분포 외(Out-of-Distribution, OOD): 학습 시 보지 못한 시퀀스 위치에서, 저주파 RoPE 성분은 학습 분포를 벗어난 회전 각도를 생성하여 불안정한 어텐션 점수를 유발합니다.
어텐션 확산(Attention Diffusion): 키(key) 세트가 커짐에 따라 소프트맥스(softmax) 분포가 평탄해져 무관한 토큰으로 확률 질량이 분산되며, 종종 중간 문맥 정보에 대한 어텐션이 부족해지는 "U자형" 위치 편향이 동반됩니다.

기존의 제로샷 방법들(예: NTK, YaRN, Self-Extend, DCA)은 일반적으로 사전에 고정된 하나의 리스케일링 계수(rescaling factor) 또는 그룹 크기를 결정합니다. 이는 트레이드오프를 강요합니다: 공격적인 계수는 짧은 문맥의 충실도를 희생시키고, 보수적인 계수는 긴 문맥에서 깔끔한 외삽(extrapolation)에 실패합니다.

방법론: Jet-Long

Jet-Long은 로컬 RoPE를 충실히 따르는 윈도우와 현재 시퀀스 길이에 따라 리스케일링 계수가 동적으로 변하는 롱 레인지(long-range) 윈도우를 결합하는 튜닝이 필요 없는 제로샷 문맥 확장 방법을 제안합니다.

1. 동적 바이포컬(Dynamic Bifocal) RoPE

핵-혁신은 고정 계수의 트레이드오프를 피하는 동적 바이포컬 매핑입니다:

로컬 윈도우(Local Window): 크기 $w_0$ 의 윈도우는 클래식 RoPE를 유지하여, 짧은 입력에 대해 베이스 모델의 사전 학습 동작을 정확하게 보존합니다.
리모트 윈도우(Remote Window): $w_0$ 를 넘어가는 위치에는 동적 위치 매핑 $f(\cdot)$ 이 적용됩니다. 연속적으로 주파수를 보간하는 방식과 달리, Jet-Long은 **이산적 위치 에일리어싱(discrete position aliasing)**을 사용합니다.
동적 그룹화(Dynamic Grouping): 그룹 크기 $G$ 는 현재 시퀀스 길이 $L$ 과 사전 학습된 윈도우 $w_{pretrained}$ 를 기반으로 동적으로 계산됩니다:
$G = \max\left(1, \left\lceil \frac{L}{w_{pretrained}} \right\rceil\right)$
리모트 매핑은 바닥 나눗셈(floor division)을 사용합니다: $f(x) = \lfloor x/G \rfloor$ . 이는 모든 리모트 회전 각도가 **분포 내(in-distribution)**에 머물게 하는 동시에, 최소한의 압축을 적용함으로써 위치 해상도를 극대화함을 보장합니다.

2. 추론 효율성 메커니즘

Jet-Long은 이 방법이 계산 효율적이며 재학습이나 무거운 KV 캐시 수정이 필요하지 않도록 두 가지 핵심 기술을 채택합니다:

온더플라이 교정 회전(On-the-Fly Correction Rotation): KV 캐시는 정확한 절대 위치에서의 압축되지 않은 베이스 키를 저장합니다. 생성 중에 Jet-Long은 캐시를 다시 쓰는 대신, 활성 쿼리와 캐시된 키에 교정 회전 $\Delta$ 를 온더플라이로 적용합니다. RoPE 회전은 가법적으로 결합되므로, 물리적 캐시에 손을 대지 않고도 레지스터 상에서 외삽된 벡터를 재구성할 수 있어, 길이에 따른 병목 없이 스트리밍 생성을 가능하게 합니다.
포함-배제 어텐션 병합(Inclusion–Exclusion Attention Merge): 프리필(prefill) 단계에서 전체 이차 어텐션 행렬을 실체화하는 것을 피하기 위해, Jet-Long은 포함-배제 원리를 사용하여 세 가지 표준 FlashAttention 패스를 결합합니다:
1. 로컬 윈도우 어텐션 (베이스 쿼리/키).
2. 전체 인과적(causal) 어텐션 (리모트 쿼리/키).
3. 로컬 윈도우 어텐션 (리모트 쿼리/키).
  LogSumExp 통계량을 통해 병합을 안정화함으로써, 이 방법은 불리언 마스크 행렬 없이 정확한 거리 기반 라우팅 출력을 계산하며 FlashAttention의 메모리 효율성을 유지합니다.

이 구성 요소들은 단일 CuTe 커널로 퓨즈되어 커널 런치 오버헤드를 최소화합니다.

주요 기여

Jet-Long 알고리즘: 리모트 윈도우 리스케일링 계수가 시퀀스 길이에 따라 동적으로 적응하여, 리모트 회전을 분포 내에 유지하면서 베이스 모델을 원래 문맥 내에서 정확하게 재현하는 튜닝 프리 바이포컬 방법입니다.
효율적인 구현: KV 캐치를 건드리지 않는 포함-배제 어텐션 병합과 온더플라이 교정 회전을 사용합니다. 이를 단일 Cuute 커널로 퓨즈함으로써, 긴 문맥에서 H100 기준 1.28–1.39× FlashAttention-2 (FA2) 프리필 처리량을 달성하고 싱글 배치 생성 시 ≤4%의 오버헤드만을 발생시킵니다.
실증적 검증: 128K 문맥까지 확장된 Qwen3-1.7B/4B/8B 모델에 대한 광범위한 평가를 통해, RULER, HELMET-RAG, PG-19 벤치마크에서 강력한 제로샷 베이스라인(DNTK, Yaarna, DCA, Self-Extend)보다 우수한 성능을 입증했습니다.
일반화: 이 방법은 하이브리드 어텐션 아키텍처(예: Jet-Nemotron)로 변경 없이 전이되며, 단일 하이퍼파라미터 $w_0$ 의 선택에 대해 견고하게 유지됩니다.

실험 결과

32K 네이티브 윈도우를 128K로 확장한 Qwen3 베이스 모델(1.7B, 4B, 8B)에 대해 평가되었습니다:

RULER 벤치마크: Jet-Long은 1.7B, 4B, 8B 모델에서 각각 **+4.79, +2.18, +2.03 퍼센티지 포인트(pp)**만큼 가장 강력한 베이스라인들을 앞섭니다. 이는 베이스 모델의 32K 학습 윈도우 내 성능을 보존하면서 확장된 길이에서 모든 베이스라인을 능가합니다.
HELMET-RAG: 롱 컨텍스트 성능의 가장 효율적인 예측 지표로 식별된 HELMET 연구에서 최고의 전체 정확도를 달성했습니다.
PG-19 퍼플렉시티(Perplexity): 모든 외삽 방법 중 가장 낮은 퍼플렉시티를 기록했습니다. 베이스 모델이 붕괴하는 동안(예: 1.7B 모델의 128K에서 퍼플렉시티 >70), Jet-Long은 낮은 퍼플렉시티(예: 1.7B의 경우 11.41)를 유지합니다.
하이브리드 아키텍처: Jet-Nemotron(소프트맥스와 선형 어텐션이 교차되는 구조)에 적용했을 때, Jet-Long은 베이스 하이브리드 모델이 붕괴하는 128K 지점에서 RULER 정확도를 +10.01 pp (2B) 및 +11.31 pp (4B) 높였습니다.
절제 연구(Ablation Studies):
- 하이퍼파라미터 탄력성: 로컬 윈도우 크기 $w_0$ 는 견고합니다. 512에서 4064 사이의 값은 최적의 성능에서 1~2 pp 이내를 유지합니다.
- 에일리어싱 vs 보간: 위치 에일리어싱(Jet-Long의 방식)이 주파수 보간(YaRN 스타일)보다 상당한 차이로 앞섭니다(예: 64K에서 +6.99 pp). 이는 회전을 이산적인 사전 학습 그리드 위에 유지하는 것의 이점을 확인시켜 줍니다.

의의

본 논문은 Jet-Long이 긴 문맥 학습의 계산 비용이나 고정 계수 외삽의 충실도 손실 없이, RoPE 레벨에서의 위치-OOD 실패 모드를 해결한다고 주장합니다. 리스케일링 계수를 동적으로 만들고 퓨즈드 커널을 통해 추론 구성을 "본질적으로 무료"로 만듦으로써, Jet-Long은 오픈 웨이트 LLM을 위한 실용적이고 고성능인 배포 경로를 제공합니다. 이는 어텐션 확산을 다루는 아키텍처적 대안(희소 어텐션 또는 상태 공간 모델 등)을 보완하며, 재학습이 필요 없는 제로샷 솔루션을 제공합니다.