Text-to-Distribution Prediction with Quantile Tokens and Neighbor Context

1. Il Problema: Regressione Testuale verso Distribuzioni

Molte applicazioni basate sui Large Language Models (LLM) richiedono la previsione di un'intera distribuzione condizionale (ad esempio, la distribuzione dei prezzi di un immobile o dei tempi di risposta a una domanda) piuttosto che un singolo valore puntuale.
Attualmente, l'approccio standard per la regressione distribuzionale con LLM presenta tre limitazioni fondamentali:

1. Mancanza di grounding locale: Le previsioni si basano solo sul testo di input, ignorando evidenze empiriche da istanze simili.
2. Collo di bottiglia delle rappresentazioni condivise: Gli approcci esistenti (come quello di Vedula et al., 2025) utilizzano uno stato nascosto finale condiviso da cui partono diverse "teste" lineari per prevedere i vari quantili. Questo crea un percorso indiretto tra l'input e l'output specifico per ogni quantile, costringendo il modello a comprimere tutte le informazioni distribuzionali in un singolo vettore.
3. Supervisione limitata: I metodi di retrieval-augmented esistenti spesso associano ai vicini recuperati solo un'etichetta puntuale, perdendo la ricchezza delle distribuzioni empiriche complete.

2. Metodologia Proposta

Gli autori introducono un nuovo framework che combina due innovazioni principali: Quantile Token Regression e Retrieval-Augmented Distribution Estimation.

A. Quantile Token Regression

Invece di utilizzare teste lineari separate su uno stato condiviso, il modello inserisce token quantile dedicati ($\langle Q_{\tau_1} \rangle, \dots, \langle Q_{\tau_Q} \rangle$) direttamente nella sequenza di input.

• Meccanismo: Questi token partecipano al meccanismo di self-attention del Transformer. Ogni token quantile può quindi "prestare attenzione" (attend) a diverse parti dell'input e accumulare informazioni specifiche per quel livello di quantile.
• Vantaggio: Crea un percorso diretto input-uscita per ogni quantile. Ad esempio, il token per il 10° percentile può focalizzarsi su segnali di "popolarità" (per risposte veloci), mentre il token per il 90° percentile può focalizzarsi su indicatori di "complessità" (per risposte lente), migliorando l'allineamento e l'interpretabilità.

B. Retrieval-Augmented Distribution Estimation

Il modello viene arricchito recuperando istanze semanticamente simili (vicini) da un pool di candidati.

• Innovazione chiave: A differenza dei metodi precedenti che recuperano solo un valore puntuale, ogni vicino recuperato viene fornito al modello insieme alla sua distribuzione empirica completa, rappresentata come un insieme di quantili osservati.
• Obiettivo: Fornire al modello evidenze distribuzionali locali per stimare non solo la tendenza centrale, ma anche la dispersione e il comportamento delle code (tail behavior) della distribuzione target.

C. Funzioni di Perdita e Analisi Teorica

Gli autori analizzano teoricamente diverse funzioni di perdita per la supervisione tramite quantili empirici:

• Wasserstein Loss ($\ell_1$ e $\ell_2$): Tratta i quantili empirici come misurazioni rumorose dei quantili della popolazione. Sono stati dimostrati essere Fisher-consistent e ottimali per l'apprendimento della forma della distribuzione.
• Pinball Loss (su quantili o mediana): Applicare il classico pinball loss ai target quantili empirici introduce un bias sistematico (inflazione/deflazione) e porta a una perdita di informazioni sulla forma della distribuzione, specialmente nelle code.

3. Contributi Chiave

1. Architettura Quantile Token: Prima implementazione che inserisce token dedicati per i quantili nella sequenza di input, permettendo percorsi diretti input-output attraverso l'attenzione.
2. Retrieval di Distribuzioni: Estensione del retrieval augmentation per includere distribuzioni empiriche complete dei vicini, non solo etichette puntuali.
3. Analisi Teorica delle Loss: Chiarificazione matematica su quali obiettivi distribuzionali ottimizzano diverse funzioni di perdita, dimostrando la superiorità delle loss di tipo Wasserstein in questo contesto specifico.
4. Benchmark Estesi: Valutazione su due dataset diversificati (Airbnb per i prezzi, StackSample per i tempi di risposta) con modelli LLM che vanno da 1.7B a 14B parametri.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli Qwen3 (1.7B - 14B) su due dataset: Inside Airbnb (prezzi) e StackSample (tempi di risposta).

• Prestazioni Superiori: L'approccio combinato (Quantile Tokens + Retrieval) supera costantemente i baseline.
  • Su StackSample (dataset più piccolo e difficile), si osserva un miglioramento drastico: riduzione del 68% dell'errore MAPE medio (da 266.65 a 84.30) rispetto al baseline senza retrieval, e un restringimento degli intervalli di previsione di un fattore 131x (RCIW@99).
  • Su Airbnb, si ottiene una riduzione del MAPE medio del 63% rispetto al baseline senza retrieval e intervalli di previsione 2 volte più stretti.
• Impatto del Retrieval: L'aggiunta di vicini recuperati migliora significativamente le prestazioni, specialmente quando i dati di addestramento sono limitati, fornendo un contesto distribuzionale ricco.
• Scalabilità: I benefici sono più marcati su dataset più piccoli e difficili. L'aumento delle dimensioni del modello (da 1.7B a 14B) porta a guadagni marginali sulla precisione puntuale, ma non garantisce automaticamente intervalli di confidenza più stretti senza l'uso delle tecniche proposte.
• Analisi delle Loss: L'uso della Wasserstein $\ell_1$ loss produce il miglior compromesso tra accuratezza e nitidezza (sharpness), mentre le loss basate su Pinball generano intervalli di previsione inutilmente ampi.

5. Significato e Implicazioni

Questo lavoro rappresenta un passo avanti significativo nella regressione testuale basata su LLM. Dimostra che:

1. La struttura interna del modello (quantile tokens) è cruciale quanto la qualità dei dati per prevedere distribuzioni complesse.
2. Il retrieval augmentation è efficace non solo per recuperare fatti, ma per groundare le stime probabilistiche confrontandole con distribuzioni empiriche di casi simili.
3. La scelta della funzione di perdita è critica: per la supervisione tramite quantili empirici, le metriche basate sulla distanza di Wasserstein sono superiori alle tradizionali loss di regressione quantile.

In sintesi, il metodo proposto offre un approccio scalabile, semplice ed efficace per trasformare il testo in distribuzioni di probabilità accurate e ben calibrate, con applicazioni potenziali in finanza, previsione della domanda, valutazione del rischio e pricing dinamico.