Jet-Long: 동적 이중초점 RoPE를 이용한 효율적인 장문맥 확장

초록

최신 LLM들은 점점 더 긴 컨텍스트를 필요로 하는 검색 증강 생성, 저장소 수준 코딩, 에이전트 워크플로우 등에 배포되고 있으며, 이러한 작업에서 누적된 추론 및 도구 사용 기록은 입력을 사전 학습 윈도우보다 한 자릿수 이상 길게 만들어 제로샷 컨텍스트 확장이 오픈웨이트 체크포인트의 주요 배포 경로가 되고 있다. 기존의 대부분의 제로샷 방법은 단일 재조정 비율을 미리 고정하는데, 공격적인 비율은 짧은 컨텍스트 충실도를 희생하고 보수적인 비율은 긴 컨텍스트에서 작동하지 않는다. 본 논문에서는 튜닝이 필요 없는 제로샷 방법인 Jet-Long을 제안한다. 이 방법은 로컬 RoPE 충실 윈도우와 현재 시퀀스 길이에 따라 동적으로 재조정 비율을 적응시키는 장거리 윈도우를 쌍으로 구성하여, 짧은 입력에서는 기본 모델을 정확히 복원하고 긴 입력에서는 깔끔하게 외삽한다. 포함-배제 어텐션 병합과 실시간 RoPE 보정 회전을 통해 이중 초점 구조는 추론 시 사실상 오버헤드가 없으며, 단일 CuTe 커널에 통합되어 긴 컨텍스트 프리필은 H100에서 FA2 처리량의 최대 1.39배(호퍼 전용 FA4에 근접)에 도달하고, 단일 배치 생성은 모든 길이에서 4% 이하의 오버헤드를 유발한다. Qwen3-1.7B/4B/8B 모델에서 최대 128K 컨텍스트까지 평가한 결과, Jet-Long은 RULER에서 가장 강력한 기준선 대비 각각 +4.79/+2.18/+2.03 퍼센트포인트의 성능 향상을 보였으며, HELMET-RAG(HELMET이 다운스트림 긴 컨텍스트 성능의 가장 효율적인 예측 변수로 식별한 벤치마크)에서 최고의 전체 정확도를 달성하고 PG-19 퍼플렉서티에서 가장 낮은 값을 기록했다. Jet-Long은 또한 재훈련 없이 긴 컨텍스트 성능을 추가로 개선하기 위해 Jet-Nemotron과 같은 하이브리드 어텐션 아키텍처로 일반화되며, 배포 용이성을 위해 하이퍼파라미터에 강건한 특성을 유지한다.

English

Modern LLMs are increasingly deployed in long-context applications such as retrieval-augmented generation, repository-level coding, and agentic workflows whose accumulated reasoning and tool traces routinely push the input an order of magnitude past the pretraining window, making zero-shot context extension the dominant deployment path for open-weight checkpoints. Most existing zero-shot methods fix a single rescaling factor up front, so an aggressive factor sacrifices short-context fidelity while a conservative one breaks down at long contexts. We propose Jet-Long, a tuning-free zero-shot method that pairs a local RoPE-faithful window with a long-range window whose rescaling factor adapts dynamically to the current sequence length, recovering the base model exactly at short inputs while extrapolating cleanly at long ones. An inclusion-exclusion attention merge and an on-the-fly RoPE correction rotation make the bifocal construction essentially free at inference; fused into a single CuTe kernel, long-context prefill reaches up to 1.39times FA2 throughput on H100 (approaching the Hopper-only FA4), and single-batch generation incurs le 4% overhead at every length. On Qwen3-1.7B/4B/8B up to 128K context, Jet-Long leads RULER by +4.79/+2.18/+2.03~pp over the strongest baseline at 1.7B/4B/8B, achieves the best overall accuracy on HELMET-RAG (a benchmark identified by HELMET as the most efficient predictor of downstream long-context performance) and attains the lowest PG-19 perplexity. Jet-Long also generalizes to hybrid attention architectures such as Jet-Nemotron for further long-context improvement without retraining, and remains hyperparameter-resilient for ease of deployment.