Pareto-Optimal Offline Reinforcement Learning via Smooth Tchebysheff Scalarization

Dit artikel introduceert STOMP, een nieuwe offline versterkingsleer-algoritme dat gebruikmaakt van gladde Tchebysheff-scalarisatie om meerdere conflicterende beloningen in één keer te optimaliseren, wat leidt tot superieure prestaties bij het afstemmen van eiwittaalmodellen op meerdere doelen vergeleken met bestaande methoden.

De kern

Stel je voor dat je een zeer getalenteerde kok bent die net een nieuw recept heeft bedacht. Je wilt dat je gerecht perfect is, maar je staat voor een lastige keuze:

• Het moet extreem lekker zijn (veel smaak).
• Maar het moet ook gezond zijn (weinig suiker en vet).
• En het moet snel klaar zijn (niet uren in de oven).

Het probleem? Je kunt niet altijd alles tegelijk maximaliseren. Als je meer suiker toevoegt voor de smaak, wordt het minder gezond. Als je het sneller wilt bakken, wordt het misschien minder gaar. In de wereld van kunstmatige intelligentie (AI) noemen we dit een meervoudig doel-probleem. De AI moet een "perfecte balans" vinden tussen deze strijdende wensen.

Deze paper introduceert een nieuwe, slimme manier om die balans te vinden, genaamd STOMP.

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het oude probleem: De "Gemiddelde" Valstrik

Vroeger probeerden AI-onderzoekers dit op een simpele manier op te lossen. Ze maakten één "super-score" door de verschillende doelen (smaak, gezondheid, snelheid) met elkaar te vermenigvuldigen en op te tellen.

• Voorbeeld: Ze zeiden: "Laten we 50% wegen op smaak en 50% op gezondheid."

Het probleem hiermee is dat dit werkt als een lineaire lijn. Het kan alleen de "gemiddelde" oplossingen vinden. Het mist echter de geheime, unieke combinaties die misschien niet perfect zijn in alles, maar juist de beste compromis zijn voor een specifieke situatie. In wiskundige termen noemen ze dit het "niet-convexe gebied" van de Pareto-voorkant. Klinkt ingewikkeld, maar stel je voor dat je een kaart hebt van alle mogelijke gerechten. De oude methode tekende alleen een rechte lijn door het midden en zei: "Hier zijn de beste opties." Maar de echte parels zaten vaak in de hoeken van de kaart, die de rechte lijn miste.

2. De nieuwe oplossing: STOMP (De Slimme Smaakmaker)

De auteurs van dit paper hebben een nieuwe methode bedacht, STOMP (Smooth Tchebysheff Optimization of Multi-Objective Preferences).

In plaats van de ingrediënten (de beloningen) direct op te tellen, kijken ze naar het hele recept als een puzzel die op een slimme manier moet worden opgelost. Ze gebruiken een wiskundige truc (genaamd "Smooth Tchebysheff scalarization") die werkt als een slimme kompasnaald.

• Hoe werkt het?
  Stel je voor dat je een kompas hebt dat niet alleen naar het noorden wijst, maar dat automatisch de sterkte van de wind en de stroming meet. Als de ene doelstelling (bijv. snelheid) al heel goed is, laat het kompas je rustig gaan. Maar als een doelstelling (bijv. gezondheid) achterblijft, duwt het kompas je hard in die richting.

  De oude methode (de rechte lijn) duwde je altijd in dezelfde hoek, ongeacht hoe goed je al was in de andere dingen. STOMP past zich dynamisch aan. Het zorgt ervoor dat je elk mogelijk compromis kunt vinden, zelfs de rare en complexe combinaties die de oude methode over het hoofd zag.

3. Waarvoor gebruiken ze dit? (De Proef)

De auteurs hebben dit getest op eiwitten (de bouwstenen van het leven).

• De taak: Ze wilden nieuwe eiwitten ontwerpen die tegelijkertijd:

  1. Zeer actief zijn (goed werk doen).
  2. Zeer specifiek zijn (niet per ongeluk andere dingen doen).
  3. Stabiel zijn (niet kapotgaan).

  Dit is als het ontwerpen van een auto die tegelijkertijd de snelste, veiligste en zuinigste ter wereld is.

Ze gebruikten drie verschillende "AI-koks" (grote taalmodellen voor eiwitten) en lieten ze leren met de nieuwe STOMP-methode. Ze vergeleken dit met de oude methoden.

4. Het resultaat

Het resultaat was indrukwekkend:

• De oude methoden vonden soms goede oplossingen, maar misten vaak de beste combinaties.
• STOMP vond in 8 van de 9 situaties de beste set van oplossingen. Het kon de "Pareto-voorkant" (de verzameling van alle mogelijke perfecte compromissen) veel vollediger afdekken.
• Het was ook robuust: het werkte goed op verschillende soorten eiwitten en met verschillende basis-modellen.

Samenvatting in één zin

Waar oude AI-methoden alleen de "gemiddelde" beste oplossing zochten en daarbij de meest interessante en unieke compromissen misten, is STOMP als een slimme navigator die je door de hele kaart van mogelijkheden leidt, zodat je altijd de perfecte balans vindt tussen strijdende doelen, of het nu gaat om het ontwerpen van medicijnen, chatbots of kunst.

Het is dus een nieuwe, krachtige tool om AI te helpen complexe, echte wereldproblemen op te lossen waar er nooit één "perfect" antwoord is, maar wel een perfect compromis.

Technische samenvatting

Probleemstelling

Grote taalmodellen (LLM's) en eiwit-taalmodellen (PLM's) worden vaak afgestemd op menselijke voorkeuren via Offline Reinforcement Learning (RL) op gelabelde datasets. Hoewel single-objective alignment (bijv. alleen "hulpvaardigheid" of "activiteit") goed onderzocht is, vereisen veel real-world toepassingen de gelijktijdige optimalisatie van meerdere, vaak conflicterende doelen. Voorbeelden zijn:

• Chatbots: Hulpvaardigheid versus veiligheid.
• Eiwitengineering: Katalytische activiteit versus specificiteit, of activiteit versus stabiliteit.

Het doel is niet één oplossing te vinden, maar de Pareto-front: de verzameling van alle niet-gedomineerde oplossingen waar geen enkel doel kan worden verbeterd zonder een ander te verslechteren.

De bestaande aanpak voor multi-objective RL is vaak lineaire reward scalarisatie (het nemen van een gewogen som van de beloningen). Het artikel benadrukt echter een fundamenteel tekortkoming: lineaire scalarisatie kan wiskundig bewezen geen oplossingen vinden in de niet-convexe gebieden van de Pareto-front. Deze gebieden vertegenwoordigen echter vaak de meest waardevolle compromissen tussen conflicterende doelen.

Methodologie: STOMP

De auteurs introduceren STOMP (Smooth Tchebysheff Optimization of Multi-Objective Preferences), een nieuw offline RL-algoritme dat de beperkingen van lineaire scalarisatie overwint.

1. Kernidee: Scalarisatie van het optimalisatieprobleem
In plaats van de rewards direct te scalariseren (zoals bij lineaire methoden), formuleren de auteurs het multi-objective RL-probleem zelf als een optimalisatieprobleem dat gescalariseerd wordt met Smooth Tchebysheff Scalarization (STS). STS is een techniek die in staat is om de volledige Pareto-front te vinden, inclusief niet-convexe delen, en is differentieerbaar (in tegenstelling tot de klassieke min-max Tchebysheff).

2. Smooth Tchebysheff Reward Scalarisatie
Een uitdaging bij STS is de gevoeligheid voor de schaal van individuele rewards. De auteurs lossen dit op door een dynamische standaardisatie toe te passen op basis van de waargenomen verdelingen in de offline dataset.

• Ze definiëren een "distribution-relative reward" $\rho_i(x, y) = r_i(x, y)/\sigma_i - \gamma \log \hat{Z}_i(x)$.
• Hierbij wordt de reward genormaliseerd door de variantie ($\sigma_i$) en gecentreerd rond een partition functie ($\hat{Z}_i$) die wordt geschat vanuit de trainingdata.
• Dit zorgt ervoor dat zeldzame, maar waardevolle sequenties in een rechts-scheve verdeling zwaarder worden bestraft dan veelvoorkomende sequenties met lage scores, wat lineaire methoden vaak verkeerd doen.

3. Het STOMP-verlies
Het algoritme bouwt voort op Direct Preference Optimization (DPO) en OffsetDPO. Het verliesfunctie ($L_{STOMP}$) combineert:

• De voorkeur voor een "winnaar" ($y_w$) boven een "verliezer" ($y_l$) gebaseerd op de gescalariseerde STS-reward.
• Een regularisatieterm die de log-likelihood-ratio tussen de huidige policy $\pi$ en een referentie-policy $\pi_0$ maximaliseert, terwijl de KL-divergentie binnen een straal wordt gehouden.
• Een extra regularisatie ($\alpha$) op de negatieve log-likelihood van de winnaars om te voorkomen dat de absolute waarschijnlijkheid van de generaties degradeert (een veelvoorkomend probleem bij offline RL).

Het algoritme gebruikt een policy-onafhankelijke versie van de reward om de paren te selecteren (voor stabiliteit), maar gebruikt de policy-afhankelijke versie in de loss-berekening voor een strakkere theoretische bound.

Belangrijkste Bijdragen

1. Nieuw Formuleringskader: De eerste toepassing van STS op het multi-objective RL-probleem zelf, in plaats van alleen op de rewards, wat leidt tot een principieel onderbouwde scalarisatie.
2. Dynamische Standaardisatie: Een methode om rewards te normaliseren op basis van hun empirische verdeling, waardoor hyperparameters voor schaling overbodig worden en de Pareto-front vollediger wordt afgedekt.
3. STOMP Algorithm: Een praktisch, offline multi-objective RL-algoritme dat direct uitbreidt op DPO en geschikt is voor taalmodellen.
4. Empirische Validatie: Uitgebreide tests op complexe biologische taken, waarbij STOMP state-of-the-art baselines (zoals DPO-Lin, ODPO-Lin en ODPO-STZ) overtreft.

Resultaten

De auteurs hebben STOMP getest op drie eiwitengineering-datasets met drie verschillende autoregressieve eiwit-taalmodellen (ProGen3-3B, ProGen-RA-3B, ProGen-RA-10B):

• DHFR: Optimalisatie van activiteit met en zonder een inhibitor (TMP). Doelen zijn ongecorreleerd.
• PbrR: Optimalisatie van lood-binding versus zink-binding. Doelen zijn sterk negatief gecorreleerd.
• $\alpha$-Amylase: Simultane optimalisatie van activiteit, expressie en thermostabiliteit. Doelen zijn positief gecorreleerd.

Kernbevindingen:

• Hypervolume: STOMP bereikte de hoogste hypervolumes (een maatstaf voor de kwaliteit en dekking van de Pareto-front) in 8 van de 9 evaluatiesettingen, zowel in offline off-policy evaluaties als in generatieve evaluaties.
• Robuustheid: In tegenstelling tot baselines die inconsistent presteerden afhankelijk van het dataset en model, was STOMP consistent superieur.
• Generatieve Kwaliteit: Bij het genereren van nieuwe eiwitsequenties (met name bij PbrR en DHFR) produceerde STOMP sequenties met de beste voorspelde eigenschappen, zelfs bij kleine steekproefgroottes.
• Niet-convexe Gebieden: De resultaten bevestigen dat STOMP in staat is om oplossingen te vinden in gebieden waar lineaire scalarisatie faalt (vooral zichtbaar in de PbrR-taak met sterk tegengestelde doelen).

Significantie

Dit werk is significant omdat het een fundamentele beperking van multi-objective RL in taalmodellen oplost: het onvermogen om niet-convexe Pareto-optima te vinden.

• Voor Eiwitengineering: Het biedt een krachtig instrument voor het ontwerpen van eiwitten met meerdere, vaak tegenstrijdige eigenschappen, wat essentieel is voor de ontwikkeling van nieuwe medicijnen en biotechnologie.
• Algemene Toepasbaarheid: Hoewel getest op eiwitten, is de methode domein-agnostisch. Het kan direct worden toegepast op andere multi-objective alignment taken, zoals het afstemmen van chatbots (hulpvaardigheid vs. veiligheid) of text-to-image generatie (kwaliteit vs. prompt-afstemming).
• Theoretische Vooruitgang: Het koppelt recent werk in multi-objective optimalisatie (STS) aan de moderne RLHF/DPO-landbouw, wat een nieuwe standaard biedt voor het afstemmen van modellen op complexe, multi-variabele doelen.

Kortom, STOMP biedt een robuust en theoretisch onderbouwd kader om taalmodellen te trainen die complexe, multi-variabele compromissen kunnen vinden, iets dat met eerdere lineaire methoden niet mogelijk was.