Sintesi Tecnica: Prototype Transformer (ProtoT)

Problematica

I modelli linguistici (LM) autoregressivi allo stato dell'arte, come GPT-4 e LLaMA, ottengono prestazioni superiori in diversi domini ma soffrono di una mancanza di trasparenza nei loro processi di ragionamento. I loro percorsi computazionali interni sono ampiamente opachi, il che porta a sfide nella comprensione di come vengano generati gli output, rischi di allucinazione e difficoltà nella modifica mirata del comportamento. Gli attuali metodi di interpretabilità sono principalmente post-hoc (ad esempio, visualizzazione dell'attenzione, probing, autoencoder sparsi) e operano su architetture che non sono state progettate con l'interpretabilità come vincolo primario. Inoltre, i meccanismi di auto-attenzione standard comportano costi computazionali quadratici rispetto alla lunghezza della sequenza, limitando la scalabilità.

Metodologia: Il Prototype Transformer (ProtoT)

Gli autori introducono ProtoT, un'architettura di LM autoregressivo che sostituisce il meccanismo di auto-attenzione standard a costo quadratico con un modulo di comunicazione basato su prototipi a costo lineare. ProtoT mantiene la struttura complessiva di LLaMA-3 (inclusi i feed-forward networks SwiGLU e l'RMS pre-layernorm) ma altera fondamentalmente lo strato "mixer".

Architettura Core

1. Mixer Basato su Prototipi: Invece dell'auto-attenzione, ProtoT utilizza $R$ vettori di parametri apprendibili chiamati prototipi. Questi prototipi agiscono come filtri che definiscono $R$ canali di comunicazione indipendenti.
2. Porte di Scrittura e Lettura (Write and Read Gates):
   • Write Gate: Aggrega le informazioni dai token passati ($x_{j}, j < i$) nei $R$ canali. Calcola i punteggi di similarità tra i token di input e i prototipi, utilizzando una softmax sui prototipi per instradare l'informazione in specifici canali. Questo crea un "collo di bottiglia di routing semantico" dove molteplici significati sono costretti a una rappresentazione di compromesso per canale, incoraggiando la specializzazione semantica.
   • Read Gate: Recupera le informazioni dai $R$ canali in base alla similarità tra il token corrente e i prototipi.
3. Prefix Mean (PM) & Time Discounting: L'informazione aggregata in ogni canale è mantenuta come media prefissa delle attivazioni passate. Un parametro di sconto temporale apprendibile ($\beta_k$) applica una Media Mobile Esponenziale (EMA) alla media prefissa, permettendo a diversi prototipi di operare a diverse scale temporali (dipendenze a breve termine vs lungo termine).
4. Complessità Lineare: L'architettura raggiunge una complessità computazionale lineare ($O(R)$) rispetto alla lunghezza della sequenza, poiché i prototipi non interagiscono tra loro (a differenza dell'interazione $O(R^2)$ negli spazi latenti stile Perceiver) e la ricorrenza permette una generazione di token a tempo costante ($O(1)$) tramite caching.
5. Alpha Gate: Una porta scalare di tipo ReZero viene applicata all'output di ogni modulo mixer per debuggare la stabilità dell'addestramento e garantire che il mixer contribuisca alle predizioni.

Contributi Chiave

1. Innovazione Architetturale: L'introduzione di ProtoT, una nuova architettura di LM che sostituisce l'auto-attenzione con una comunicazione multi-canale basata su prototipi, raggiungendo una complessità computazionale lineare pur mantenendo prestazioni competitive.
2. Interpretabilità Emergente: Dimostrazione che i prototipi apprendono automaticamente concetti disgiunti e nominabili (ad esempio, "donna", "Nuova Zelanda", "COVID-19") durante l'addestramento attraverso vari livelli di astrazione. Ciò avviene senza supervisione esplicita, guidato dalla pressione architettonica a specializzare i canali.
3. Editing Comportamentale Mirato: Evidenza che il ragionamento del modello può essere editato chirurgicamente intervenendo su specifici prototipi (ad esempio, ri-inizializzando o mascherando un prototipo "femmina"). Tali interventi alterano significativamente la probabilità di token correlati (ad esempio, diminuendo "donne" di circa il 16%) causando una degradazione trascurabile della perplexity complessiva del modello (<1% nella maggior parte dei casi).
4. Analisi della Robustezza: Dimostrazione del fatto che ProtoT mostra robustezza rispetto a perturbazioni lessicali che preservano il significato (sinonimi, errori di battitura) e che questa robustezza è mediata dai meccanismi di routing dei prototipi.

Risultati Sperimentali

Gli autori hanno valutato ProtoT contro tre baseline: LLaMA (auto-attenzione standard), Mamba (modello a spazio di stato) e DeltaNet (attenzione lineare), utilizzando un backbone fisso (6 strati, 256 dimensione nascosta) e addestrando su 250 milioni di token di FineWeb-Edu.

• Scalabilità: ProtoT scala bene con la dimensione del modello e il volume dei dati. In configurazioni su larga scala (512 dimensione nascosta, 12 strati), ProtoT raggiunge una perplexity di 29.5, superando DeltaNet (31.5) e Mamba (26.5), sebbene resti dietro LLaMA (25.8).
• Generazione di Testo: Nelle valutazioni di preferenza umana a coppie (LLM-as-a-judge), ProtoT supera LLaMA e DeltaNet, posizionandosi vicino a Mamba.
• Task Downstream: Nel benchmark GLUE, ProtoT rimane competitivo con le baseline dense, spesso eguagliando o superando Mamba e DeltaNet, e dimostrando una forte robustezza cross-domain.
• Throughput: ProtoT scala meglio di LLaMA all'aumentare della lunghezza del contesto grazie alla sua computazione lineare, superando LLaMA a 32k token, sebbene DeltaNet mantenga il throughput più elevato a contesti molto lunghi.
• Metriche di Interpretabilità: I prototipi di ProtoT ottengono punteggi più alti in termini di disgiunzione (disentanglement) e copertura dei temi principali rispetto alle teste di attenzione di LLaMA e alle feature degli Autoencoder Sparsi (SAE).
• Efficacia dell'Intervento: La perturbazione di specifici prototipi (ad esempio, L9 P7 per "femmina") ha causato cali significativi nelle probabilità dei token target (ad esempio, -16.60% per "donne") con danni collaterali minimi alla coerenza della sequenza.

Significato e Claim

Il paper sostiene che ProtoT rappresenti un passo promettente verso modelli linguistici autoregressivi interpretabili per design. Integrando l'interpretabilità direttamente nell'architettura tramite canali di prototipi, il modello:

• Cattura concetti coerenti e nominabili automaticamente.
• Permette editi mirati e chirurgici al comportamento del modello senza destabilizzare l'intero sistema.
• Offre un'alternativa alla computazione lineare all'auto-attenzione che è robusta alle perturbazioni.

Gli autori riconoscono i limiti, notando che, sebbene le prestazioni siano competitive con le baseline lineari, attualmente sono inferiori ai modelli di auto-attenzione stabiliti come LLaMA in termini di perplexity pura. Essi avvertono inoltre che le etichette dei prototipi dovrebbero essere trattate come approssimazioni piuttosto che come verità assolute, poiché alcuni prototipi rimangono polisemantici. Tuttavia, il lavoro suggerisce che bilanciare capacità e trasparenza è fattibile, offrendo una via per modelli linguistici più sicuri e controllabili.