Technische Samenvatting: Het Leren van Metrische Afstand aan Discrete Autoregressieve Taalmodellen

Probleemstelling

Grote Taalmodellen (LLMs) en andere discrete autoregressieve fundamentele modellen werken door tokens uit een vocabulaire te voorspellen, waarbij elke token wordt behandeld als een discrete categorische variabele die wordt geoptimaliseerd tegen een one-hot doeldistributie. Hoewel deze formulering de aanpassing van LLMs aan diverse domeinen zoals visie, robotica en multimodaal redeneren mogelijk heeft gemaakt, negeert deze fundamenteel de metrische relaties tussen tokens.

In taken die numeriek of metrisch gestructureerde elementen omvatten—zoals gehele getallen in rekenkundig redeneren, ruimtelijke coördinaten in objectdetectie, rotatiehoeken in robotmanipulatie of gekwantiseerde embeddings in beeldgeneratie—gaat de intrinsieke afstandsstructuur verloren wanneer tokens worden gereduceerd tot one-hot doelen. Conventioneel trainen negeert dat een voorspelling van "4" semantisch dichter bij een ground truth van "5" ligt dan een voorspelling van "2". Deze beperking hindert de data-efficiëntie en prestaties in domeinen waar afstand betekenisvol is.

Methodologie: DIST2Loss

De auteurs introduceren DIST2Loss (DIScreTized DISTance Loss), een afstandsbewuste doelfunctie voor discrete autoregressieve modellen. De kernidee is om de standaard one-hot doeldistributie te vervangen door een beloningsgewogen distributie die is afgeleid van vooraf gedefinieerde tokenafstanden.

Kernformulering

Gegeven een vocabulairesubset $V_d$ en een metrische ruimte $(X, d)$, construeert DIST2Loss een doeldistributie $p_d(v|x, t)$ voor een token $v$ op tijdstap $t$, gegeven een ground-truth token $x$. Deze distributie wordt gedefinieerd met behulp van een gediscretiseerde exponentiële familie:

$$p_d(v|x, t) = \frac{\exp\left(-\frac{d(v, x, t)}{\tau}\right)}{\sum_{v' \in V_d} \exp\left(-\frac{d(v', x, t)}{\tau}\right)}$$

Waarbij:

• $d(v, x, t)$ de afstand is tussen de kandidaat-token $v$ en de ground truth $x$ (bijvoorbeeld Euclidische afstand voor coördinaten, hoekafstand voor rotaties, of embeddende afstand voor VQ-codes).
• $\tau$ een temperatuur-hyperparameter is die de gladheid van de distributie controleert.

Het uiteindelijke trainingsdoel combineert de standaard cross-entropy loss ($L_{CE}$) met de afstandsbewuste regularisatie:
$$L = L_{CE} + \alpha L_{dist}$$
waarbij $L_{dist}$ de KL-divergentie is tussen de door het model voorspelde distributie en de afstandsbewuste doel $p_d$, en $\alpha$ de twee termen in evenwicht brengt.

Theoretische Interpretatie

DIST2Loss wordt geïnterpreteerd als de gesloten-vorm oplossing voor entropie-geregulariseerde beleidsoptimalisatie. In versterkend leren (RL) moedigt entropie-regularisatie exploratie aan door te voorkomen dat het beleid te vroeg instort. Het optimale beleid voor het maximaliseren van de verwachte beloning met entropie-regularisatie staat bekend als evenredig te zijn met $\exp(R(a)/\tau)$. DIST2Loss instantieert dit door kandidaat-tokens te behandelen als acties en de negatieve afstand als een beloningssignaal. In tegenstelling tot traditioneel RL, dat bemonstering, rollouts en iteratieve updates vereist, biedt DIST2Loss een stabiele, variantievrije, analytische doeldistributie, waarmee de instabiliteit die geassocieerd is met beleidsgradiëntschatters wordt vermeden.

Omgaan met Hoge Dimensies en Multi-Token Sequenties

Het kader is flexibel genoeg om hoge-dimensionale continue vectoren (bijvoorbeeld VQ-embeddings) te verwerken door afstandsmaatstaven te definiëren over hun embeddings (bijvoorbeeld cosijn-sequentie). Voor multi-token gestructureerde elementen (zoals meercijferige getallen) merken de auteurs op dat het direct toepassen van de loss een "credit assignment-probleem" oplevert, waarbij het model niet kan onderscheiden welke token een fout veroorzaakte. Om dit te mitigeren, stellen ze praktische oplossingen voor zoals Contrastive Target Augmentation (bemonsteren van negatieve buren) en Place Value Weighting (loss wegen op basis van positionele significantie), hoewel hun primaire experimenten zich richten op taken waar beloningen natuurlijk decomponeren op token-niveau.

Belangrijkste Bijdragen

1. DIST2Loss Kader: Een plug-and-play doel dat metrische structuur integreert in discrete autoregressieve training zonder extra data of aanzienlijke rekenkosten te vereisen.
2. Theoretische Brug: Het leggen van een verbinding tussen afstandsbewuste supervisie en entropie-geregulariseerd RL, en het bieden van een stabiel alternatief voor versterkend leren voor taken met bekende per-token beloningen.
3. Generalisatie: Het aantonen dat de methode toepasbaar is in diverse domeinen, van 1D-scalar regressie tot hoge-dimensionale beeldgeneratie, mits de tokens een betekenisvolle metrie toelaten.

Experimentele Resultaten

De auteurs hebben DIST2Loss geëvalueerd over vijf verschillende domeinen, waarbij consistent verbeteringen werden aangetoond in data-efficiëntie en downstream prestaties:

1. Synthetische Functieregressie (Toy Task): In een meta-learning-opstelling waar modellen lineaire regressie leren uit slechts 1–10 voorbeelden, presteerde DIST2Loss significant beter dan standaard SFT en vocabulaire-beperkte baselines in termen van Mean Absolute Error (MAE) en Root Mean Square Error (RMSE), wat superieure structurele generalisatie aantoont.
2. Visuele Grounding: Bij het fine-tunen van Phi3V op RefCOCO/RefCOCO+/RefCOCOg zonder voorafgaande objectdetectietraining, bereikte DIST2Loss prestaties die vergelijkbaar waren met state-of-the-art modellen die zijn voorgetraind op grote schaal detectiedatasets. Opmerkelijk is dat zelfs in "moeilijke gevallen" waarbij beide modellen het verkeerde object voorspelden, de DIST2Loss-voorspellingen geometrisch dichter bij de ground truth lagen (hogere IoU).
3. Robotmanipulatie: Op de VIMABench verbeterde DIST2Loss de succespercentages van robotactiegeneratie (gewelhoeken en coördinaten). De winsten waren het meest uitgesproken in regimes met een tekort aan data (1K voorbeelden), wat de waarde van metrische priors benadrukt wanneer data beperkt is.
4. Generatieve Beloningsmodellering: Toegepast op LLM-uitlijning (RLHF), verbeterde DIST2Loss de prestaties van beloningsmodellering op RewardBench en MT-Bench. De generatieve aanpak met DIST2Loss presteerde beter dan zowel standaard binaire classifiers als SFT-baselines, met name in de domeinen "Chat Hard" en "Safety".
5. Beeldgeneratie: Met behulp van vector-gekwantiseerde (VQ) representaties, verbeterde DIST2Loss het leren van beeldrepresentaties in autoregressieve modellen (LlamaGen), wat resulteerde in lagere FID en hogere IS-scores in vergelijking met SFT-baselines over verschillende modelgroottes en trainingsepochen.

Betekenis en Claims

Het artikel beweert dat afstandsbewuste supervisie een eenvoudig en algemeen alternatief biedt voor one-hot supervisie voor discrete autoregressieve modellen. De betekenis ligt in het vermogen om:

• Intrinsieke Structuur te Benutten: De metrische relaties die inherent zijn aan veel niet-tekstuele taken (coördinaten, getallen, embeddins) te exploiteren, die anders worden genegeerd door standaard cross-entropy.
• Data-efficiëntie te Verbeteren: Betere prestaties te behalen met minder trainingsvoorbeelden door inductieve bias direct in de lossfunctie te coderen.
• Training te Stabiliseren: Een gesloten-vorm, stabiel alternatief te bieden voor versterkend leren voor beloningsuitlijning, waarmee de bemonsterings- en instabiliteitsproblemen van beleidsgradiënten worden vermeden.

De auteurs benadrukken dat de methode het meest effectief is wanneer tokens een betekenisvolle metrie toelaten (bijvoorbeeld numerieke waarden, coördinaten, gekwantiseerde embeddins). In domeinen zonder intrinsieke geometrie reduceert de methode van nature tot standaard one-hot supervisie. Het werk suggereert dat voor een breed scala aan multimodale en redeneertaken het integreren van metrische afstand een cruciaal, maar tot nu toe onderbenut, onderdeel is van het trainen van discrete fundamentele modellen.