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

Sintesi Tecnica: Jet-Long

Definizione del Problema

I moderni Large Language Models (LLM) sono sempre più impiegati in applicazioni a contesto lungo, come la generazione aumentata dal recupero (RAG), la programmazione a livello di repository e i workflow agentici. In questi scenari, l'accumulo di tracce di ragionamento e output di strumenti spinge spesso le lunghezze di input di un ordine di grandezza oltre la finestra di pre-addestramento (ad esempio, 100K+ token rispetto a una finestra di addestramento di 32K).

L'addestramento diretto dei modelli su contesti lunghi rimane proibitivamente costoso a causa dei costi computazionali quadratici e della scarsità di dati a lungo contesto, mentre il fine-tuning spesso degrada le prestazioni a contesto breve. Di conseguenza, l'estensione del contesto zero-shot (estendere il contesto senza fine-tuning) è diventata la via di deployment dominante per i checkpoint a pesi aperti.

Tuttavia, i modelli LLM vanilla basati su Transformer che utilizzano l'Embedding Posizionale Rotatorio (RoPE) falliscono nel generalizzare oltre la loro finestra di addestramento a causa di due modalità di fallimento primarie:

1. Fuori Distribizione (OOD) della Posizione: Alle posizioni di sequenza non viste durante l'addestramento, le componenti RoPE a bassa frequenza producono angoli di rotazione al di fuori della distribuzione di addestramento, causando punteggi di attenzione erratici.
2. Diffusione dell'Attenzione: Man mano che l'insieme delle chiavi cresce, la distribuzione softmax si appiattisce, disperdendo la massa di probabilità su token irrilevanti, spesso accompagnata da un bias posizionale a "forma di U" che sottostima l'informazione nel contesto centrale.

I metodi zero-shot esistenti (ad esempio, NTK, YaRN, Self-Extend, DCA) tipicamente fissano un singolo fattore di riscalamento o una dimensione di raggruppamento in anticipo. Ciò impone un compromesso: un fattore aggressivo sacrifica la fedeltà a contesto breve, mentre uno conservativo non riesce a estrapolare pulitamente ai contesti lunghi.

Metodologia: Jet-Long

Jet-Long propone un metodo di estensione del contesto zero-shot e senza tuning che accoppia una finestra fedele al RoPE locale con una finestra a lungo raggio il cui fattore di riscalamento si adatta dinamicamente alla lunghezza corrente della sequenza.

1. RoPE Bifocale Dinamico

L'innovazione principale è una mappatura bifocale dinamica che evita il compromesso del fattore fisso:

• Finestra Locale: Una finestra di dimensione $w_0$ mantiene il RoPE classico, preservando esattamente il comportamento di pre-addestramento del modello base per gli input brevi.
• Finestra Remota: Per le posizioni oltre $w_0$, viene applicata una mappatura della posizione dinamica $f(\cdot)$. A differenza dei metodi che interpolano le frequenze continuamente, Jet-Long utilizza l'aliasing posizionale discreto.
• Raggruppamento Dinamico: La dimensione del gruppo $G$ è calcolata dinamicamente in base alla lunghezza corrente della sequenza $L$ e alla finestra pre-addestrata $w_{pretrained}$:
  $$G = \max\left(1, \left\lceil \frac{L}{w_{pretrained}} \right\rceil\right)$$
  La mappatura remota è una divisione intera (floor division): $f(x) = \lfloor x/G \rfloor$. Ciò assicura che ogni angolo di rotazione remoto rimanga in distribuzione (sulla griglia di pre-addestramento del modello) massimizzando al contempo la risoluzione posizionale applicando la compressione minima necessaria.

2. Meccanismi di Efficienza dell'Inferenza

Per garantire che il metodo sia computazionalmente efficiente e non richieda il riaddestramento o pesanti modifiche alla cache KV, Jet-Long impiega due tecniche chiave:

• Rotazione di Correzione On-the-Fly: La cache KV memorizza le chiavi base non compresse nelle loro esatte posizioni assolute. Durante la generazione, invece di riscrivere la cache, Jet-Long applica una rotazione di correzione $\Delta$ alla query attiva e alle chiavi memorizzate on-the-fly. Poiché le rotazioni RoPE si compongono additivamente, questo ricostruisce i vettori estrapolati nei registri senza toccare la cache fisica, permettendo la generazione in streaming attraverso i confini di lunghezza senza blocchi.
• Merge dell'Attenzione Inclusione-Esclusione: Per evitare la materializzazione dell'intera matrice di attenzione quadratica durante il prefill, Jet-Long combina tre passaggi standard di FlashAttention utilizzando un principio di inclusione-esclusione:
  1. Attenzione della finestra locale (query/chiavi base).
  2. Attenzione causale completa (query/chiavi remote).
  3. Attenzione della finestra locale (query/chiavi remote).
     Stabilizzando il merge con statistiche LogSumExp, il metodo calcola l'esatto output di routing basato sulla distanza senza una matrice di maschera booleana, mantenendo l'efficienza di memoria di FlashAttention.

Questi componenti sono fusi in un singolo kernel CuTe, minimizzando l'overhead di lancio del kernel.

Contributi Chiave

1. Algoritmo Jet-Long: Un metodo bifocale senza tuning dove il fattore di riscalamento della finestra remota si adatta dinamicamente alla lunghezza della sequenza, mantenendo le rotazioni remote in distribuzione e riproducendo esattamente il modello base all'interno del suo contesto nativo.
2. Implementazione Efficiente: Un merge dell'attenzione inclusione-esclusione e una rotazione di correzione on-the-fly che lasciano intatta la cache KV. Fusi in un singolo kernel CuTe, questi raggiungono un throughput di prefill di 1.28–1.39× rispetto a FlashAttention-2 (FA2) su H100 ai contesti lunghi e comportano un overhead $\le$ 4% durante la generazione a singolo batch.
3. Validazione Empirica: Valutazione estesa sui modelli Qwen3-1.7B/4B/8B fino a 128K di contesto, dimostrando prestazioni superiori rispetto ai forti baseline zero-shot (DNTK, YaRN, DCA, Self-Extend) sui benchmark RULER, HELMET-RAG e PG-19.
4. Generalizzazione: Il metodo si trasferisce invariato ad architetture di attenzione ibrida (ad esempio, Jet-Nemotron) e rimane robusto alla scelta dell'unico iperparametro $w_0$.

Risultati Sperimentali

Valutato sui modelli base Qwen3 (1.7B, 4B, 8B) con una finestra nativa di 32K estesa a 128K:

• Benchmark RULER: Jet-Long guida i più forti baseline con +4.79, +2.18 e +2.03 punti percentuali (pp) rispettivamente sui modelli 1.7B, 4B e 8B. Preserva le prestazioni del modello base all'interno della finestra di addestramento di 32K e supera tutti i baseline ai contesti estesi.
• HELMET-RAG: Raggiunge la migliore accuratezza complessiva su questo benchmark, identificato dallo studio HELMET come il predittore più efficiente delle prestazioni a lungo contesto a valle.
• Perplessità PG-19: Ottiene la perplessità più bassa tra tutti i metodi di estrapolazione. Mentre il modello base collassa (ad esempio, perplessità >70 a 128K per 1.7B), Jet-Long mantiene una perplessità bassa (ad esempio, 11.41 per 1.7B).
• Architetture Ibride: Applicato a Jet-Nemotron (interleaving di softmax e linear attention), Jet-Long eleva l'accuratezza RULER di +10.01 pp (2B) e +11.31 pp (4B) a 128K, dove il modello ibrido base altrimenti collasserebbe.
• Studi di Ablazione:
  • Resilienza dell'Iperparametro: La dimensione della finestra locale $w_0$ è robusta; valori compresi tra 512 e 4096 rimangono entro 1–2 pp dalle migliori prestazioni.
  • Aliasing vs. Interpolazione: L'aliasing posizionale (l'approccio di Jet-Long) supera l'interpolazione di frequenza (stile YaRN) con margini significativi (ad esempio, +6.99 pp a 64K), confermando il beneficio di mantenere le rotazioni sulla griglia discreta pre-addestrata.

Significato

Il paper sostiene che Jet-Long affronti il fallimento del modo di errore Position-OOD a livello di RoPE senza il costo computazionale dell'addestramento a lungo contesto o la perdita di fedeltà dell'estrapolazione a fattore fisso. Rendendo il fattore di riscalamento dinamico e la costruzione dell'inferenza "essenzialmente gratuita" tramite kernel fusi, Jet-Long fornisce una via di deployment pratica e ad alte prestazioni per i modelli LLM a pesi aperti in scenari a lungo contesto. Completa le alternative architettoniche (come l'attenzione sparsa o i modelli state-space) che affrontano la diffusione dell'attenzione, offrendo una soluzione zero-shot che non richiede riaddestramento e si generalizza attraverso diverse dimensioni dei modelli e architetture ibride.