OLLM: Options-based Large Language Models

1. Il Problema

I Large Language Models (LLM) tradizionali operano prevedendo una singola distribuzione per il prossimo token, affidandosi a euristiche di campionamento (come temperatura, top-k o top-p) per introdurre diversità nelle generazioni. Tuttavia, questo approccio presenta diverse limitazioni:

• Inefficienza nel controllo: La ricerca di strategie ottimali avviene su uno spazio di azione enorme (l'intero vocabolario), rendendo l'ottimizzazione tramite Reinforcement Learning (RL) costosa in termini di campioni e instabile.
• Allineamento e deriva: L'addestramento con RLHF (Reinforcement Learning from Human Feedback) può portare a "reward hacking" o a comportamenti indesiderati, come il cambio di lingua o ragionamenti degenerati, poiché l'esplorazione non è vincolata alle modalità apprese durante il fine-tuning supervisionato (SFT).
• Multimodalità nascosta: In molti contesti (specialmente nel ragionamento matematico), esistono multiple continuazioni plausibili per un dato token. I modelli standard tendono a fondere queste opzioni in una singola distribuzione, spesso favorendo una modalità dominante e perdendo percorsi di ragionamento validi alternativi.

2. Metodologia: OLLM (Options-based LLM)

Gli autori propongono OLLM, un metodo generale che sostituisce la previsione del singolo prossimo token con un insieme di opzioni apprese, indicizzate da una variabile latente discreta.

Architettura

OLLM agisce come un "plug-in" leggero che si inserisce prima del layer di output (lm_head) di qualsiasi LLM pre-addestrato, aggiungendo solo due moduli:

1. Encoder ($E_\psi$): Durante l'addestramento supervisionato, prende l'embedding del token target ($y_{t+1}$) e lo stato nascosto contestuale ($h_t$) per produrre una distribuzione su uno spazio latente discreto $Z = \{1, ..., K\}$. Questo permette di mappare un token specifico a diverse "opzioni" a seconda del contesto (es. diverse traiettorie di ragionamento).
2. Decoder ($D_\omega$): Prende lo stato nascosto $h_t$ e un'opzione latente $z$ campionata dalla distribuzione dell'encoder, generando un vettore di bias ($\Delta h_t$) che viene aggiunto allo stato nascosto prima della proiezione sul vocabolario.

Meccanismo di Funzionamento

• Decomposizione della Distribuzione: Invece di competere per la massa di probabilità all'interno di un singolo softmax, i token alternativi vengono ottimizzati sotto diverse latenti. Questo disaccoppia le continuazioni concorrenti.
• Policy nel Spazio Latente: Durante l'inferenza, l'encoder viene sostituito da una policy leggera ($\pi_\phi$) che, dato lo stato nascosto $h_t$, seleziona l'opzione latente $z$ appropriata.
• Vantaggio Chiave: La policy opera in uno spazio di azione a bassa dimensionalità (dimensione $K$, es. 10) invece che sull'intero vocabolario. Questo vincola l'esplorazione alle opzioni apprese durante l'SFT, riducendo drasticamente la varianza e prevenendo allucinazioni o deviazioni semantiche.

Obiettivi di Addestramento

Il modello viene addestrato con una combinazione di tre perdite:

1. Cross-Entropy ($L_{CE}$): Per garantire che il token corretto venga predetto data l'opzione latente.
2. KL Divergence ($L_{KL}$): Una regolarizzazione adattiva verso una distribuzione uniforme per prevenire il collasso delle latenti (latent collapse) e garantire che lo spazio delle opzioni sia ben sfruttato.
3. Loss della Policy ($L_{policy}$): Una policy viene addestrata per imitare l'encoder (behavioral cloning), prevedendo la latente corretta direttamente dallo stato nascosto.

3. Risultati Sperimentali

Gli esperimenti sono stati condotti su un modello backbone di 1.7B parametri (una versione distillata di Qwen2.5/Deepseek R1), addestrato sul dataset OpenMathReasoning e valutato su OmniMath. Solo il 1.56% dei parametri è stato reso addestrabile.

• Performance di Base: I modelli baseline adattati con LoRA (State-of-the-Art) raggiungono un picco di 51% di correttezza nella risposta finale.
• Performance OLLM:
  • Con la selezione ottimale della latente (oracolo), OLLM raggiunge circa il 70% di correttezza nella risposta finale e l'83% di accuratezza per token.
  • Questo dimostra che lo spazio categorico delle opzioni è in grado di catturare efficacemente le variazioni nel testo che i modelli standard non riescono a mantenere senza sovrapposizione.
• Efficienza: L'uso di uno spazio latente compatto rende l'ottimizzazione della policy molto più efficiente in termini di campioni rispetto al RL tradizionale sul vocabolario completo.
• Robustezza: OLLM riduce significativamente le disallineamenti comuni (come il cambio di lingua o ragionamenti degenerati) poiché la policy è vincolata alle opzioni apprese durante l'SFT.

4. Contributi Chiave

1. Modellazione Esplicita delle Opzioni: Introduzione di un metodo che espone esplicitamente le opzioni per token attraverso variabili latenti discrete, invece di affidarsi al campionamento stocastico.
2. Architettura Plug-in Leggera: Un approccio che richiede l'aggiunta di solo due layer (encoder/decoder) e un numero minimo di parametri addestrabili, rendendo la conversione di qualsiasi LLM pre-addestrato in un OLLM semplice e scalabile.
3. Policy Learning nello Spazio Latente: Dimostrazione che l'addestramento di una policy in uno spazio latente a bassa dimensionalità migliora l'efficienza del RL e la stabilità dell'allineamento, senza bisogno di penalità KL esplicite o loss di allineamento artigianali.
4. Miglioramento nel Ragionamento Matematico: Evidenza empirica che la separazione delle modalità di ragionamento in diverse latenti permette di preservare percorsi di soluzione validi che altrimenti verrebbero soppressi.

5. Significato e Prospettive Future

Il lavoro di OLLM segna un passo importante verso LLM più controllabili, robusti ed efficienti.

• Interpretabilità: Poiché diverse latenti possono corrispondere a diverse traiettorie di ragionamento, il modello offre una maggiore trasparenza su come vengono prese le decisioni.
• Futuro del RL: L'approccio suggerisce che il Reinforcement Learning per gli LLM dovrebbe essere condotto in spazi latenti compatti e ben strutturati piuttosto che sullo spazio del vocabolario grezzo, per evitare problemi di esplorazione e allineamento.
• Applicazioni: Oltre al ragionamento matematico, gli autori propongono di estendere OLLM alla generazione di codice, QA scientifica e pianificazione multi-step, dove la capacità di esplorare percorsi logici alternativi è cruciale.

In sintesi, OLLM trasforma la generazione di testo da un processo di selezione monodimensionale in un processo di gestione di opzioni strutturate, offrendo un nuovo paradigma per migliorare le capacità di ragionamento dei modelli linguistici.