SWE-Pruner Pro: The Coder LLM Already Knows What to Prune

Riepilogo Tecnico: SWE-Pruner Pro

Definizione del Problema

Gli agenti di codifica che operano in ambienti multi-turn accumulano estesi output di strumenti (ad es. cat, grep, ls, log di esecuzione python) durante la loro interazione con i codebase. Questi output contengono spesso una significativa ridondanza, consumando la maggior parte del budget di token per traiettoria e innescando il degrado del contesto lungo. Le soluzioni di pruning del contesto esistenti si dividono in due categorie:

1. Compressione generalista: Metodi come LLMLingua2 o Selective Context assegnano punteggi ai token basandosi su metriche fisse (ad es. perplessità, struttura sintattica). Questi non riescono ad adattarsi all'evoluzione del focus del task dell'agente.
2. Pruning specifico per il task: Il lavoro precedente come SWE-Pruner condiziona il pruning sull'intento dell'agente, ma richiede un modello di scoring esterno e una query esplicita di "goal-hint" che l'agente deve generare ad ogni turno. Ciò introduce overhead e latenza aggiuntivi.

Entrambi gli approcci trattano il segnale di pruning come una quantità esterna da ricostruire, ignorando il fatto che il modello backbone dell'agente ha già elaborato l'output dello strumento e probabilmente ha codificato le informazioni di rilevanza all'interno delle sue rappresentazioni interne.

Metodologia: SWE-Pruner Pro

SWE-Pruner Pro propone di leggere il segnale di pruning direttamente dalle rappresentazioni interne (hidden states) del backbone dell'agente, eliminando la necessità di modelli di scoring esterni o esplicite query di "goal-hint". Il metodo opera direttamente all'interno del loop di inferenza dell'agente.

Meccanismo Core

1. Utilizzo della Rappresentazione Interna: Quando l'agente legge una risposta dello strumento ($r_t$), il backbone congelato esegue un prefill standard, generando gli stati nascosti dell'ultimo layer ($h_i$) per ogni token. SWE-Pruner Pro attacca una testa di pruning leggera a questi stati esistenti.
2. Architettura della Testa di Pruning:
   • Embedding Sensibile alla Lunghezza: Un embedding appreso $e(N)$, condizionato sul numero totale di righe $N$ della risposta, viene aggiunto tramite broadcast a ogni stato nascosto. Ciò consente al modello di regolare la strategia di pruning in base alla lunghezza della risposta (ad es. essendo più conservativo con output brevi dove ogni riga è critica).
   • Classificatore Non Lineare: Una piccola rete feed-forward (LayerNorm + due blocchi Linear-GELU-Dropout) mappa gli stati nascosti aumentati ai logit di "keep-or-prune" (mantieni o scarta). Questa non linearità è necessaria per risolvere la sovrapposizione nello spazio di rappresentazione osservata negli esperimenti di probing lineare.
   • Aggregazione delle Decisioni: I logit a livello di token vengono aggregati tramite voto di maggioranza all'interno di ogni riga per produrre una decisione binaria di mantenimento/scarto per l'intera riga.
3. Strategia di Training:
   • Dati: La testa è addestrata su 22.609 traiettorie multi-turn annotate da Claude Sonnet 4.6 con etichette per riga di tipo keep/prune.
   • Funzione di Perdita: Gli autori impiegano una per-sample balanced focal loss. A differenza della focal loss standard che bilancia globalmente, questo approccio calcola la perdita separatamente per i token "keep" e "prune" all'interno di ogni campione e li media con lo stesso peso. Ciò assicura che i campioni con tassi di mantenimento estremi (ad es. solo il 3% delle righe mantenute) contribuiscano equamente all'obiettivo, evitando che il modello si adatti eccessivamente al tasso di mantenimento medio (~30%).
4. Integrazione nell'Inferenza: La testa viene eseguita "in-server" durante la fase di prefill. Essa riutilizza gli stati nascosti generati per la risposta dello strumento, aggiungendo solo un singolo passaggio forward della piccola testa. La risposta potata ($\tilde{r}_t$) sostituisce la risposta grezza nel contesto per il turno successivo, mentre la generazione del turno corrente attende ancora alla risposta grezza completa.

Contributi Chiave

1. Scoperta del Segnale: Il paper dimostra che le rappresentazioni interne del backbone di un agente codificano già l'importanza a livello di riga degli output degli strumenti. Un semplice probe lineare raggiunge un AUC di 0.83, provando che il segnale "keep-or-prune" esiste senza necessità di ricostruzione esterna.
2. Architettura: Introduzione di una testa leggera, SWE-Pruner Pro, che utilizza un embedding sensibile alla lunghezza e una per-sample balanced focal loss per prendere decisioni di pruning precise direttamente dagli stati congelati del backbone.
3. Efficienza: Il metodo evita chiamate supplementari al modello e query esplicite di "goal-hint", aggiungendo un overhead di inferenza limitato pur ottenendo una significativa riduzione dei token.

Risultati Sperimentali

Il metodo è stato valutato su due backbone open-weight (MiMo-V2-Flash e Qwen3-Coder-Next) su quattro benchmark multi-turn: SWE-Bench Verified, SWE-QA, SWE-QA-Pro e Oolong.

• Risparmio di Token: SWE-Pruner Pro è l'unico metodo tra i sette valutati che riduce l'uso end-to-end dei token in ogni configurazione.
  • Fino al 39% di riduzione su SWE-QA-Pro (Qwen3-Coder-Next).
  • 30% di riduzione sul benchmark a lungo contesto Oolong (MiMo-V2-Flash).
  • Al contrario, altri metodi (ad es. LLMLingua2) talvolta hanno aumentato l'uso dei token fino al 190% su task a lungo contesto a causa dell'overhead.
• Qualità del Task:
  • SWE-QA/Pro: Mantiene la qualità del task entro una banda stretta rispetto alla baseline non potata (ad es. variazione del punteggio di +0.02 su SWE-QA).
  • Oolong: Migliora l'accuratezza di +2.2 punti su MiMo-V2-Flash.
  • SWE-Bench Verified: Su MiMo-V2-Flash, migliora il tasso di risoluzione di +3.8% riducendo al contempo i token di input del 7.4% rispetto alla baseline non potata. Su Qwen3-Coder-Next, ottiene la maggiore riduzione dei token di input (-13.5%) con la minima degradazione del tasso di risoluzione (-1.2%).
• Latenza: La testa in-server aggiunge circa il 15.0% del tempo wall-clock aggregato al tempo di generazione totale, un costo compensato dalla riduzione della dimensione del contesto per i turni successivi.

Significato e Rivendicazioni

Il paper sostiene che il "giudizio di rilevanza" richiesto per il pruning del contesto non è un segnale esterno che deve essere elicitato tramite query o modelli aggiuntivi, ma è già presente nelle rappresentazioni formate dal backbone mentre legge passivamente un'osservazione.

• Cambio di Paradigma: Invece di costruire il pruning intorno all'agente (richiedendo modelli esterni o query esplicite), SWE-Pruner Pro semplicemente legge il segnale che l'agente ha già formato.
• Impatto Pratico: Questo approccio risolve il compromesso tra qualità e compressione osservato nel lavoro precedente, permettendo un pruning aggressivo (fino al 39% di risparmio di token) senza degradare le prestazioni del task e, in alcuni casi, migliorandole mitigando il degrado del lungo contesto.
• Generalizzabilità: Sebbene valutato su task di codifica incentrati su Python, il pipeline è indipendente dal linguaggio ed è stato dimostrato che preserva il risparmio di token e la qualità anche su task a lungo contesto in linguaggio naturale (Oolong), suggerendo che il paradigma si trasferisce tra diversi domini.

Gli autori concludono che, per gli agenti di codifica, la via più efficiente per la gestione del contesto è sfruttare lo stato interno del backbone stesso, evitando la latenza e la complessità dei meccanismi di scoring esterni.