StateBridge: Training-free Hidden-state Alignment for Latent Communication in LLM Multi-Agent Systems

Riepilogo Tecnico: StateBridge

Problematica

I sistemi Multi-Agente (MA) basati su Large Language Models (LLM) attualmente si affidano al linguaggio naturale (token discreti) per la comunicazione inter-agente. Questo approccio introduce un "collo di bottiglia discreto": il mittente deve comprimere i propri stati latenti continui in token discreti, scartando informazioni che le sole identità dei token non possono catturare. Quando il ricevente mappa questi token nuovamente in vettori tramite il proprio strato di embedding, recupera solo le identità dei token, non gli stati latenti continui originali del mittente. Ciò risulta in una perdita di informazioni e in una coordinazione compromessa.

Sebbene la trasmissione diretta degli stati latenti sia un'alternativa potenziale, essa affronta un disallineamento geometrico. Gli LLM pre-addestrati sono istruiti per elaborare gli embedding dei token al livello di input, mentre gli stati latenti del decoder occupano una regione differente dello spazio di rappresentazione. Il passaggio diretto degli stati latenti causa un disallineamento semantico, poiché il ricevente non è in grado di interpretare vettori che mancano della specifica allineamento geometrico dello spazio di embedding del proprio input. Le soluzioni esistenti o iniettano gli stati interni strato per strato (limitandone la portabilità) o richiedono proiettori addestrati (limitandone la generalizzabilità a modelli specifici).

Metodologia: StateBridge

Gli autori propongono StateBridge, un'interfaccia di comunicazione priva di addestramento che allinea gli stati latenti dell'ultimo strato del mittente allo spazio di embedding di input del ricevente tramite una trasformazione ortogonale in forma chiusa. Il metodo opera in tre fasi:

1. Estrazione del Messaggio:
   L'agente mittente genera un messaggio. Invece di trasmettere l'intera sequenza, il sistema estrae gli ultimi $K$ stati latenti ($S$) dei token generati. Recupera inoltre i corrispondenti embedding dei token di riferimento ($R$) dalla matrice di embedding del vocabolario condiviso. Questi embedding di riferimento fungono da target per l'allineamento.

2. Interfaccia di Allineamento:
   Il cuore di StateBridge è una trasformazione in tre fasi per convertire $S$ in un prefisso allineato $\bar{S}$ compatibile con il ricevente:

   • Allineamento Procruste Ortogonale: Il metodo tratta $S$ e $R$ come nuvole di punti in $\mathbb{R}^d$. Per prima cosa, centra e sbianca (whitening) entrambi gli insiemi per rimuovere gli offset globali ed equalizzare le direzioni principali. Successivamente, risolve un problema di Procrustes ortogonale per trovare la matrice di rotazione ottimale $Q^*$ che mappa gli stati del mittente sbiancati verso gli embedding di riferimento sbiancati, preservando distanze e angoli tra coppie. Ciò assicura che la struttura geometrica (similarità semantica) della rappresentazione del mittente sia mantenuta.
   • Calibrazione della Norma: Gli stati latenti dell'ultimo strato hanno tipicamente norme significativamente più grandi degli embedding dei token di input. StateBridge riscala ogni vettore allineato per far corrispondere la norma $\ell_2$ media agli embedding del vocabolario. Ciò evita che il prefisso domini i punteggi di attenzione del ricevente semplicemente a causa della magnitudo.
   • Ancoraggio al Vocabolario: Per garantire che i vettori allineati si trovino nelle regioni dello spazio di embedding incontrate durante il pre-addestramento, ogni vettore calibrato viene spostato leggermente verso il suo embedding di vocabolario più vicino (basandosi sulla similarità del coseno). Questo passaggio mantiene la continuità pur ancorando la rappresentazione alla distribuzione di input attesa dal ricevente.

3. Iniezione del Prefisso:
   Il prefisso allineato risultante $\bar{S}$ viene preposto agli embedding del prompt del ricevente. Il ricevente elabora questo prefisso continuo utilizzando i propri strati standard di modellazione del linguaggio senza alcuna modifica architettonica o tokenizzazione.

Contributi Chiave

1. Risoluzione del Disallineamento di Rappresentazione: Il documento dimostra che l'allineamento in forma chiusa è sufficiente per risolvere il disallineamento geometrico tra gli stati latenti dell'ultimo strato e gli embedding di input, consentendo una comunicazione latente efficace senza addestramento o modifiche architettoniche.
2. Framework StateBridge: Gli autori propongono un'interfaccia specifica e priva di parametri che combina l'allineamento di Procrustes, la calibrazione della norma e l'ancoraggio al vocabolario. Questo approccio preserva le relazioni geometriche tra coppie degli stati del mittente garantendo al contempo la compatibilità con lo spazio di input del ricevente.
3. Superiorità Empirica: Attraverso quattro modelli di due famiglie (Qwen3 e OLMo3) e otto benchmark (matematica, QA, codice), StateBridge supera costantemente la comunicazione basata su testo e i baseline di trasferimento della KV-cache. Raggiunge il punteggio migliore o il pare al meglio in 22 su 26 coppie modello-compito.

Risultati

• Performance: StateBridge ha ottenuto un miglioramento medio del 2,4 - 2,9 punti percentuali rispetto al miglior baseline in tutte le configurazioni di modello.
• Generalizzabilità: Il metodo ha mostrato una robusta performance attraverso diverse architetture di modelli. In particolare, sul modello OLMo3-7B-Think, il trasferimento della KV-cache (LatentMAS) ha performato scarsamente (media 55,1%) rispetto al testo (73,9%), mentre StateBridge ha raggiunto il 76,7%. Ciò suggerisce che l'operazione di solo input-layer di StateBridge è più robusta alle differenze architettoniche rispetto ai metodi che iniettano stati attraverso tutti i layer del transformer.
• Sensibilità al Compito: I guadagni sono stati più pronunciati nei benchmark impegnativi che richiedono un ragionamento profondo (es. MedQA, GPQA, AIME, generazione di codice), dove la comunicazione testuale è più soggetta a scartare utili informazioni continue.
• Studi di Ablazione:
  • Sostituire l'allineamento di Procrustes con la regressione ridge (una mappa lineare non ortogonale) ha causato un calo significativo delle prestazioni (da un 82,4% a un 74,9% medio), confermando che preservare la struttura geometrica tra coppie è critico.
  • Rimuovere la calibrazione della norma o l'ancoraggio al vocabolario ha degradato anche le prestazioni, indicando che l'allineamento geometrico da solo è insufficiente senza compatibilità di scala e distribuzione.
• Caso di Studio: Un compito di ricostruzione ha dimostrato che il prefisso allineato contiene informazioni semantiche (es. termini medici specifici come "koilonychia") che erano assenti dai token del suffisso visibili, provando che il prefisso trasporta informazioni oltre l'output testuale discreto.

Significato e Rivendicazioni

Il documento sostiene che la sfida principale nella comunicazione latente non è semplicemente trasmettere stati più ricchi, ma rendere tali stati leggibili per l'agente successivo. StateBridge affronta questo problema garantendo che gli stati trasmessi siano sia informativi che geometricamente compatibili con lo spazio di input del ricevente.

Gli autori sottolineano che il loro approccio è privo di addestramento e non richiede modifiche architettoniche, rendendolo applicabile a LLM pronti all'uso. Sostengono che la progettazione dell'interfaccia di comunicazione sia un fattore critico e spesso trascurato nei sistemi MA omogenei. Sebbene il presente lavoro assuma pesi del modello condivisi e lunghezze del prefisso fisse, i risultati suggeriscono che le interfacce basate sull'allineamento possono migliorare significativamente la collaborazione multi-agente, in particolare in compiti in cui la tokenizzazione discreta causa un sostanziale smarrimento di informazioni.