Provably avoiding over-optimization in Direct Preference Optimization without knowing the data distribution

Probleemstelling

Direct Preference Optimization (DPO) is een standaard geworden voor het post-trainen van grote taalmodellen (LLM's) vanwege zijn eenvoud en het vermogen om expliciete beloningsmodellering te omzeilen. DPO staat echter bekend om over-optimalisatie: naarmate het trainingsproces vordert, kan het model de trainingsverlies minimaliseren terwijl de werkelijke kwaliteit van gegenereerde antwoorden ten opzichte van het initiële model verslechtert. Dit probleem wordt verergerd door het feit dat standaard DPO en Reinforcement Learning from Human Feedback (RLHF) de geschatte beloning behandelen als perfecte grondwaarheid, waarbij onzekerheid wordt genegeerd.

Theoretische garanties voor het mitigeren van over-optimalisatie vertrouwen doorgaans op pessimisme, waarbij wordt gewaarborgd dat het geleerde beleid robuust is tegen schattingsfouten. Bestaande pessimistische varianten (zoals $\chi^2$-geregulariseerde DPO en DPO+SFT) bereiken sterke theoretische grenzen die uitsluitend afhankelijk zijn van een concentratiecoëfficiënt voor één beleid ($C^\star$), wat robuustheid impliceert zelfs als de dataset niet het volledige beleidsruimte dekt. Deze garanties zijn echter kritiek afhankelijk van het kennen van de data-genererende verdeling ($\pi_{data}$) of het hebben van toegang tot een referentiebeleid dat goed wordt gedekt door $\pi_{data}$. In de praktijk wordt deze aanname vaak geschonden, bijvoorbeeld bij het finetunen van kleinere modellen op voorkeursdata gegenereerd door grotere, proprietaire systemen (zoals GPT-4), waarbij $\pi_{data}$ onbekend en ontoegankelijk is.

De centrale onderzoeksvraag die wordt aangepakt is: Kunnen we een DPO-achtig algoritme ontwerpen dat bewezen robuust is tegen over-optimalisatie zonder kennis van de data-genererende verdeling $\pi_{data}$?

Methodologie: PEPO

De auteurs introduceren PEPO (Pessimistic Ensemble based Preference Optimization), een single-step algoritme dat pessimisme bereikt via een ensemble van beleidsmodellen getraind op disjuncte datasubsets, zonder kennis van $\pi_{data}$ te vereisen.

1. Ensemble Training:
De voorkeursdataset $D$ wordt gepartitioneerd in $L$ disjuncte subsets. Voor elke subset $D_\ell$ wordt een lichtgewicht beleid $\tilde{\pi}_\ell$ getraind met behulp van een pessimistische DPO-doelfunctie. Deze doelfunctie modificeert de standaard DPO-loss door de sigmoid-functie horizontaal te verschuiven met een parameter $\lambda_\ell$, wat effectief de marge verkleint die nodig is voor een winnende waarschijnlijkheid. Deze verschuiving induceert een conservatieve (pessimistische) schatting van de beloningskloof.

2. Pessimistische Aggregatie:
Het uitvoerbeleid $\pi_{out}$ wordt gedefinieerd als de oplossing voor een "virtueel" RLHF-probleem dat de worst-case beloning over het ensemble optimaliseert. De impliciete beloning voor het ensemble wordt geconstrueerd door het minimum te nemen over de $L$ modellen van hun geschatte beloningskloven, aangepast voor de data-verdeling in de andere subsets (om onafhankelijkheid te waarborgen).
$$r^-(x, a) = \beta \min_{\ell \in [L]} \left( \log \frac{\tilde{\pi}_\ell(a|x)}{\pi_{ref}(a|x)} - \zeta_\ell(x) \right)$$
waarbij $\zeta_\ell(x)$ rekening houdt met de bijdrage van data buiten de $\ell$-de subset.

3. Efficiënt Sampling:
De gesloten-vorm oplossing voor $\pi_{out}$ omvat een normalisatieconstante over de volledige responsruimte, wat onberekenbaar is. Om dit te omzeilen, maakt PEPO gebruik van rejection sampling (Algoritme 2). Elk lid van het ensemble $\tilde{\pi}_\ell$ dient als voorstelverdeling, omdat de teller van $\pi_{out}$ begrensd is door $\tilde{\pi}_\ell$. Voor praktische grootschalige toepassingen stellen de auteurs ook een token-level benadering voor die de pessimistische aggregatie toepast bij elke decodestap, waardoor rejection sampling op trajectniveau wordt vermeden en de generatielatentie aanzienlijk wordt verminderd.

Belangrijkste Bijdragen

Theoretische Bijdragen:

• Staalcomplexiteit zonder $\pi_{data}$: In de tabulaire setting bereikt PEPO staalcomplexiteitsgaranties die uitsluitend afhankelijk zijn van de concentratiecoëfficiënt voor één beleid ($C^\star$). Dit staat in contrast met standaard DPO, dat lijdt aan een concentratiecoëfficiënt voor alle beleidsmodellen ($C_{all}$), en bestaande pessimistische DPO-varianten die kennis van $\pi_{data}$ vereisen.
• Optimale Ensemblegrootte: De analyse karakteriseert de optimale ensemblegrootte $L$ als de kleinste waarde die voldoende is om pessimisme met hoge waarschijnlijkheid te garanderen. Theoretische grenzen tonen aan dat $L$ logaritmisch schaalt met de grootte van de staat- en actieruimte.
• Robuustheid: De grenzen gelden zonder aan te nemen dat de dataset de volledige beleidsklasse dekt, waarmee de "onbekende $\pi_{data}$"-kloof in eerdere theoretische werken wordt aangepakt.

Praktische Bijdragen:

• Empirische Validatie: PEPO wordt geëvalueerd op modellen met 7B, 8B en 34B parameters (Zephyr-7B, Llama-3.1-8B, Mistral-7B, Yi-34B) met behulp van de UltraFeedback-dataset.
• Mitigatie van Over-optimalisatie: Experimenten tonen aan dat terwijl standaard DPO een scherpe daling in winnende percentages vertoont na initiële winsten (over-optimalisatie), de prestaties van PEPO stabiliseren en verbeteren.
• Superioriteit ten opzichte van Baselines: PEPO presteert consistent beter dan $\chi^2$PO en DPO+SFT in settings waarbij $\pi_{data} \neq \pi_{ref}$ (bijvoorbeeld UltraFeedback), waarbij de theoretische garanties van die baselines ongeldig zijn.
• Efficiëntie: De token-level implementatie bereikt een gunstige afweging tussen winnend percentage en generatietijd, en biedt bijna optimale prestaties met aanzienlijk lagere latentie dan volledige rejection sampling.

Resultaten

• Over-optimalisatie: Figuur 1 toont dat het winnende percentage van DPO ten opzichte van het initiële model piekt en vervolgens daalt, terwijl PEPO (met een geschikte $L$) een aanhoudende verbetering laat zien. Vroegtijdig stoppen op basis van validatieverlies blijkt onbetrouwbaar te zijn voor het selecteren van het beste DPO-checkpoint.
• Vergelijking: In Figuur 3 bereikt PEPO hogere winnende percentages op AlpacaEval2 in vergelijking met $\chi^2$PO en SFT+DPO over alle geteste modellen. De prestatiekloof neemt alleen af wanneer $\pi_{data} = \pi_{ref}$, wat bevestigt dat het voordeel van PEPO voortkomt uit het vermogen om onbekende data-verdelingen te hanteren.
• Ablatie: Experimenten bevestigen dat $L$ een kritieke hyperparameter is; een te kleine waarde faalt om voldoende pessimisme te induceren, terwijl een te grote waarde de effectieve data per model vermindert, wat de prestaties verslechtert.

Betekenis en Claims

Het artikel claimt de eerste DPO-achtige algoritme te presenteren dat bewezen over-optimalisatie vermijdt zonder kennis van de data-genererende verdeling. Dit is significant omdat het de kloof overbrugt tussen theoretische robuustheid (concentratie voor één beleid) en praktische beperkingen (onbekende $\pi_{data}$), een veelvoorkomend scenario in de uitlijning van LLM's in de echte wereld waar data wordt gegenereerd door proprietaire of diverse modellen.

De auteurs positioneren PEPO als een praktische oplossing die de eenvoud van DPO behoudt (single-step training, geen expliciet beloningsmodel) terwijl het de theoretische veiligheid van pessimistische regularisatie integreert. Zij erkennen dat hun theoretische garanties momenteel zijn vastgesteld in de tabulaire setting, en dat het uitbreiden van deze resultaten naar algemene functbenaderingsklassen zonder kennis van $\pi_{data}$ een open uitdaging blijft. Bovendien levert de huidige methode een niet-parametrisch beleid op (via sampling), wat ideaal is voor de laatste stap van een trainingspijplijn maar niet voor tussentijdse checkpoints, wat een richting suggereert voor toekomstig werk aan parametrische outputs.