TAB-PO: Preference Optimization with a Token-Level Adaptive Barrier for Token-Critical Structured Generation

1. Il Problema: Ottimizzazione delle Preferenze in Contesti Strutturati Critici

Il lavoro affronta una sfida specifica nell'allineamento dei Large Language Models (LLM): l'applicazione dell'Ottimizzazione Diretta delle Preferenze (DPO) a compiti di predizione strutturata token-critica, come l'annotazione medica.

Sebbene il DPO standard abbia ottenuto risultati eccellenti in compiti di generazione libera (es. riassunti, conversazioni), mostra fragilità in scenari dove:

1. Bassa separazione delle preferenze: Le coppie di risposte "preferite" (chosen) e "rifiutate" (rejected) differiscono spesso per un numero minimo di token (1-3 token), come un'etichetta gerarchica errata o un confine di span impreciso.
2. Squilibrio nell'importanza dei token: I dati strutturati (es. JSON) contengono molti token strutturali ad alta frequenza (scaffolding, chiavi JSON) che hanno poco valore semantico, mentre pochi token semantici critici (etichette gerarchiche, span di evidenza) sono essenziali per la correttezza del compito.

In questo regime, il DPO standard soffre di tre problemi principali:

• Collasso del margine: La separazione log-probabilità è insufficiente tra risposte quasi identiche.
• Compressione della verosimiglianza (Likelihood Squeezing): L'obiettivo del margine può aumentare la probabilità relativa della risposta preferita ma diminuire la sua probabilità assoluta, rendendo la generazione instabile.
• Diluizione del gradiente: Il peso uniforme a livello di sequenza disperde il segnale di apprendimento sui token strutturali comuni, lasciando i token critici (rari e confondibili) con aggiornamenti deboli e rumorosi.

2. Metodologia: TAB-PO e il Task PV-Miner

Per risolvere questi problemi, gli autori introducono TAB-PO (Token-Adaptive Barrier Preference Optimization) e un nuovo benchmark chiamato PV-Miner.

Il Task PV-Miner

Il benchmark formalizza l'annotazione della comunicazione paziente-fornitore come un compito di predizione strutturata. Il modello deve estrarre tuple gerarchiche (Codice, Sottocodice, Span) da messaggi sicuri.

• Sfide: Squilibrio delle classi, sovrapposizione di categorie comunicative, e la necessità di un ancoraggio esatto a livello di caratteri (Span).
• Dataset: Costruito su 1137 messaggi clinici reali, con annotazioni esperte che coprono 8 Codici principali e 26 Sottocodici.

L'Approccio TAB-PO

TAB-PO è un metodo di ottimizzazione post-SFT (Supervised Fine-Tuning) che modifica il DPO standard attraverso due meccanismi chiave:

1. Decomposizione e Pesatura a Livello di Token (Token-Weighted Advantage):

   • Invece di trattare l'intera sequenza come un'unità atomica, TAB-PO decompone il vantaggio log-likelihood a livello di token.
   • Assegna pesi maggiori ($w_{Code}, w_{Sub-code}, w_{Span}$) ai token semantici critici e pesi nulli o bassi ai token strutturali (JSON scaffolding).
   • Questo concentra il segnale di apprendimento sui token che determinano la correttezza, mitigando la diluizione del gradiente.

2. Barriera Adattiva a Livello di Token (Adaptive Token-Level Barrier):

   • Per prevenire la "compressione della verosimiglianza" e mantenere la validità dello schema, viene introdotta una barriera condizionale.
   • La barriera si attiva solo quando il modello è poco confidente su un token preferito (probabilità inferiore a una soglia $\tau$).
   • In tali casi, applica una penalità che "ancora" la probabilità del token alla verosimiglianza appresa durante l'SFT, regolarizzando selettivamente solo i token a rischio di deriva senza vincolare quelli già sicuri.

Costruzione delle Coppie di Preferenza

Le coppie di preferenze non sono generate casualmente, ma sono consapevoli della confusione (Confusion-Aware):

• 40% Curate da esperti: Basate su casi ambigui reali identificati durante le sessioni di revisione.
• 60% Sintetiche: Derivate dagli errori empirici del modello SFT su un set di validazione (es. confusione tra Codici/Sottocodici, annotazioni mancanti o extra).
• Le risposte rifiutate ($Y^-$) sono perturbazioni minime delle risposte corrette ($Y^+$) che mantengono la validità dello schema ma introducono errori realistici (es. cambio di sottocodice, span errato).

3. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli di diverse dimensioni (da 1.5B a 70B parametri, inclusi Llama-3 e Qwen2.5) sul task PV-Miner.

• Performance Superiori: TAB-PO supera costantemente l'SFT e le varianti DPO più recenti (IPO, SimPO, Cal-DPO, DPO-Positive).
  • Miglioramento medio relativo di ~4% in Micro-F1 rispetto all'SFT.
  • Miglioramenti particolarmente significativi nella previsione dei Sottocodici (fino a +8.31 punti F1 per Qwen2.5-1.5B), dove la distinzione tra etichette confondibili è critica.
• Stabilità: TAB-PO riduce la varianza tra le diverse esecuzioni (run-to-run variance), indicando un'ottimizzazione più stabile.
• Ablazione: Lo studio di ablazione conferma che la pesatura dei token è il componente più impattante (+1.7 punti F1), seguito dalla normalizzazione della lunghezza e dal bilanciamento delle classi.
• Confronto con Baseline: Mentre il prompting zero-shot e few-shot fallisce nella struttura e nel recall, e l'SFT migliora ma lascia errori fini, TAB-PO risolve efficacemente le ambiguità a bassa separazione.

4. Contributi Chiave

1. PV-Miner Benchmark: Un nuovo dataset e task per l'annotazione strutturata della comunicazione medica, che cattura sfide reali come l'etichettatura gerarchica, lo squilibrio delle classi e l'ancoraggio esatto degli span.
2. TAB-PO (Metodo): Un nuovo algoritmo di ottimizzazione delle preferenze che:
   • Evidenzia i token semantici critici tramite pesatura adattiva per evitare la diluizione del gradiente.
   • Introduce una barriera adattiva basata sulla confidenza per stabilizzare la generazione e prevenire il collasso della verosimiglianza.
3. Analisi delle Coppie di Preferenza: Una metodologia per costruire coppie di preferenze realistiche e a bassa separazione basate su errori empirici del modello e curate da esperti, essenziali per addestrare modelli su compiti strutturati critici.

5. Significato e Implicazioni

Questo lavoro è significativo perché sposta il paradigma dell'allineamento dei modelli linguistici dalla generazione libera a compiti strutturati ad alto rischio (come la medicina). Dimostra che le tecniche standard di DPO, progettate per margini ampi e generici, sono inadeguate quando la correttezza dipende da pochi token specifici.

TAB-PO offre una soluzione pratica per:

• Migliorare l'affidabilità degli LLM in contesti clinici dove un errore di etichettatura può avere conseguenze reali.
• Fornire un framework per l'ottimizzazione delle preferenze in scenari con dati sbilanciati e strutture rigide (JSON, XML).
• Stabilire che la gestione esplicita dell'importanza dei token e la stabilizzazione della verosimiglianza sono cruciali per il successo dell'addestramento su compiti di estrazione strutturata fine-grained.

In sintesi, TAB-PO rappresenta un passo avanti fondamentale per rendere gli LLM più robusti e affidabili nell'analisi di dati sanitari strutturati, superando i limiti delle attuali tecniche di allineamento.