Doladění a RLHF: jak se z doplňovače textu stane asistent

Nechť Σ je konečná množina znaků (tokenů), |Σ| = K, a Σ* množina všech konečných řetězců nad Σ. Jazykový model s parametry θ je zobrazení, které každému textu x ∈ Σ* přiřadí rozdělení pravděpodobnosti π_θ(· | x) na Σ, tj. π_θ(v | x) ≥ 0 pro všechna v ∈ Σ a ∑_{v ∈ Σ} π_θ(v | x) = 1. V řeči posilovaného učení je to stochastická strategie: stav je dosavadní text, akce další znak.

Pro odpověď y = y₁…y_m ∈ Σ^m, m ≥ 1, označme y_<t = y₁…y_t−1 (pro t = 1 prázdný řetězec) a definujme pravděpodobnost celé odpovědi jako součin

π_θ(y | x) = ∏_t=1^m π_θ(y_t | x, y_<t),   tedy   ln π_θ(y | x) = ∑_t=1^m ln π_θ(y_t | x, y_<t).

Nechť model po textu vrátí logity z = (z₁, …, z_K) ∈ ℝ^K (jeden výstup na každý znak) a τ > 0. Losování s teplotou τ vybere k-tý znak s pravděpodobností

p_k = exp(z_k/τ)/∑_j=1^K exp(z_j/τ),   k = 1, …, K.

Při τ = 1 je to přesně rozdělení modelu π_θ (softmax logitů) — vidíme, co model „opravdu“ říká, i s jeho nejistotou. Odpověď na dotaz x losujeme takto:

1. y := prázdný řetězec;
2. vylosuj znak v ∈ Σ s pravděpodobnostmi p_k výše (logity modelu po textu xy) a připoj ho: y := yv;
3. je-li v ∈ {„.“, „?“} nebo má y 16 znaků, skonči a vrať y; jinak pokračuj krokem 2.

Podíl druhu odpovědi je počet takových odpovědí/počet měst · 40, na šesti ladených městech tedy se jmenovatelem 240: třeba 178/240 = 74,2 %.

Nechť

1. π_θ0 je předtrénovaný (základní) jazykový model s parametry θ₀, viz díl Předtrénování;
2. 𝒟_SFT = {(x⁽ⁱ⁾, y⁽ⁱ⁾) : i = 1, …, M} je konečná množina ukázek: x⁽ⁱ⁾ ∈ Σ* je pokyn nebo otázka a y⁽ⁱ⁾ ∈ Σ* vzorová odpověď, kterou napsal člověk;
3. pro ukázku (x, y) je w = w₁…w_T = xy spojení otázky a odpovědi a m_t ∈ {0, 1} maska: m_t = 1, právě když token w_t patří k odpovědi y.

Pro dávku B ⊆ 𝒟_SFT je ztráta SFT

L_SFT(θ) = −1/|B| ∑_{(x, y) ∈ B} ∑_t=1^T m_t · ln π_θ(w_t | w_<t) = −1/|B| ∑_{(x, y) ∈ B} ln π_θ(y | x).

Doladění s učitelem (supervised fine-tuning, SFT) je minimalizace L_SFT gradientním sestupem po náhodných dávkách, se startem v θ = θ₀.

Slovy: předtrénovaný model se dál učí stejnou křížovou entropií dalšího tokenu, ale na malé sadě pečlivě připravených ukázek dialogu, a ztráta se počítá pouze na tokenech odpovědi. Učí se kratší dobu a s menším krokem, aby model nezapomněl, co umí z předtrénování.

Nechť π_ref je jazykový model po SFT (zmrazený), β > 0 a

𝒟 = {(x_i, y_i⁺, y_i⁻) : i = 1, …, M}

data preferencí: x_i ∈ Σ* je dotaz, y_i⁺, y_i⁻ ∈ Σ* jsou dvě odpovědi modelu na něj a člověk vybral y_i⁺ jako lepší (píšeme y_i⁺ ≻ y_i⁻). Označme σ(u) = 1/1 + e^−u logistickou funkci. RLHF (reinforcement learning from human feedback) je postup o dvou krocích:

1. Model odměny. Funkce r_φ : Σ* × Σ* → ℝ s parametry φ (obvykle jazykový model s jedním číselným výstupem) se naučí minimalizací L_R(φ) = −1/M ∑_i=1^M ln σ(r_φ(x_i, y_i⁺) − r_φ(x_i, y_i⁻)). To je metoda maximální věrohodnosti pro Bradleyho–Terryho model P(y⁺ ≻ y⁻ | x) = σ(r(x, y⁺) − r(x, y⁻)).
2. Posilované učení. Strategie π_θ (start v parametrech π_ref) se metodou gradientu strategie, nejčastěji PPO (díl Policy gradient a PPO), učí maximalizovat J(θ) = E_{x, y ∼ πθ(· | x)}[ r_φ(x, y) ] − β · E_x[ KL(π_θ(· | x) ‖ π_ref(· | x)) ], kde x probíhá dotazy z dat a KL je Kullbackova–Leiblerova divergence (definice níže). Epizoda má jediný krok: akce je celá odpověď y a odměna na jejím konci je r_φ(x, y) − β ln π_θ(y | x)/π_ref(y | x); z ní PPO počítá výhody Â.

Slovy: lidé dostanou dvě odpovědi modelu na stejný dotaz a vyberou lepší. Z tisíců takových porovnání se naučí model odměny, který odpovědím přiřadí skóre, a jazykový model se pak učí psát odpovědi s vysokým skóre — s trestem β · KL za vzdálení od modelu po SFT.

(Barvy podle vzorce výše.)

Nechť Y je konečná nebo spočetná množina (u nás odpovědí) a π, π_ref : Y → [0, 1] jsou rozdělení pravděpodobnosti na Y (∑_{y ∈ Y} π(y) = ∑_{y ∈ Y} π_ref(y) = 1). Předpokládejme, že π_ref(y) > 0 pro každé y s π(y) > 0. Kullbackova–Leiblerova divergence π od π_ref je

KL(π ‖ π_ref) = ∑_{y ∈ Y, π(y) > 0} π(y) · ln π(y)/π_ref(y) = E_{y ∼ π}[ ln π(y) − ln π_ref(y) ].

(Sčítance s π(y) = 0 se vynechávají, podle konvence 0 · ln 0 = 0. Existuje-li y s π(y) > 0 a π_ref(y) = 0, klade se KL(π ‖ π_ref) = +∞.)

Platí KL(π ‖ π_ref) ≥ 0 a KL(π ‖ π_ref) = 0, právě když π = π_ref (Gibbsova nerovnost). Slovy: KL měří, o kolik je v průměru přes odpovědi modelu π logaritmus jeho pravděpodobnosti vyšší než u π_ref. Není symetrická: průměruje se přes odpovědi, které dává π, takže trestá hlavně to, když π dá velkou pravděpodobnost odpovědi, kterou π_ref skoro nikdy neříká.

Nechť π_ref je zmrazený jazykový model po SFT, β > 0, 𝒟 = {(x_i, y_i⁺, y_i⁻) : i = 1, …, M} data preferencí jako v definici RLHF a σ logistická funkce. Pro strategii π_θ a trojici (x, y⁺, y⁻) označme

s_θ(x, y⁺, y⁻) = β · ( ln π_θ(y⁺ | x)/π_ref(y⁺ | x) − ln π_θ(y⁻ | x)/π_ref(y⁻ | x) ).

DPO (Direct Preference Optimization; Rafael Rafailov a kol., 2023) je minimalizace ztráty

L_DPO(θ) = −1/M ∑_i=1^M ln σ( s_θ(x_i, y_i⁺, y_i⁻) )

gradientním sestupem po náhodných dávkách, se startem v π_θ = π_ref.

Slovy: DPO řeší stejnou úlohu jako RLHF s trestem KL, ale bez modelu odměny a bez posilovaného učení. Podle věty níže optimální model π* splňuje r(x, y) = β ln π*(y | x)/π_ref(y | x) + (člen nezávislý na y), takže výraz s_θ hraje roli rozdílu odměn v Bradleyho–Terryho modelu a stačí rovnou maximalizovat věrohodnost preferencí.

Nechť x je pevný dotaz, Y konečná neprázdná množina odpovědí, π_ref rozdělení pravděpodobnosti na Y s π_ref(y) > 0 pro všechna y ∈ Y, r : Y → ℝ odměna a β > 0. Pro každé rozdělení π na Y položme

J(π) = ∑_{y ∈ Y} π(y) r(y) − β · KL(π ‖ π_ref)

a definujme

π*(y) = π_ref(y) · exp(r(y)/β)/Z,   Z = ∑_{y ∈ Y} π_ref(y) · exp(r(y)/β).

Pak π* je rozdělení na Y a pro každé rozdělení π ≠ π* na Y platí J(π) < J(π*). Maximum J přes všechna rozdělení na Y se tedy nabývá právě v π*.

Důkaz. Z > 0, π*(y) > 0 a ∑_y π*(y) = Z/Z = 1, takže π* je rozdělení. Z definice π* je π_ref(y) · exp(r(y)/β) = Z π*(y) a r(y) = β ln exp(r(y)/β). Sčítáme-li přes y s π(y) > 0, dostaneme

J(π) = −β ∑ π(y) ln π(y)/π_ref(y) · exp(r(y)/β) = −β ∑ π(y) ln π(y)/Z π*(y) = −β · KL(π ‖ π*) + β ln Z.

Druhý člen na π nezávisí a podle Gibbsovy nerovnosti je KL(π ‖ π*) ≥ 0 s rovností, právě když π = π*. ∎

Zlogaritmováním definice π* a úpravou: r(y) = β ln π*(y)/π_ref(y) + β ln Z. Dosadíme do Bradleyho–Terryho modelu — v rozdílu r(y⁺) − r(y⁻) se β ln Z odečte, protože nezávisí na odpovědi. Odměna je tak vyjádřená jen modelem samým a model odměny už není potřeba. To je celé odvození DPO.

Model si po pár krocích zvýšil ln-pravděpodobnost lepší odpovědi „česko.“ o 0,3 a horší „polsko.“ snížil o 0,2 (vůči π_ref):

(Barvy podle vzorce L_DPO výše.) Gradient ztráty je −β σ(−s) · [∇ ln π_θ(y⁺ | x) − ∇ ln π_θ(y⁻ | x)]: krok zvyšuje ln π_θ lepší odpovědi a snižuje horší, a to tím méně, čím lépe už je dvojice seřazená. Po tomhle pokroku je tah na dvojici asi poloviční.

Nechť r^*(x, y) je skutečná kvalita odpovědi y na dotaz x (to, co lidé doopravdy chtějí; přímo ji změřit neumíme) a r_φ(x, y) její náhrada, podle které se optimalizuje — model odměny, u DPO odměna skrytá v datech preferencí. Nechť θ₀, θ₁, θ₂, … jsou parametry strategie π_θ v průběhu optimalizace a

R_φ(k) = E_{x, y ∼ πθk}[ r_φ(x, y) ],   R^*(k) = E_{x, y ∼ πθk}[ r^*(x, y) ].

Říkáme, že mezi kroky k₀ < k₁ dochází k přeoptimalizaci odměny, právě když

1. R_φ(k₁) > R_φ(k₀) (podle náhradní odměny se model zlepšil) a zároveň
2. R^*(k₁) < R^*(k₀) (ve skutečnosti se zhoršil).

Slovy: model najde odpovědi, které podle odměny (nebo podle preferencí) vycházejí dobře, ale ve skutečnosti dobré nejsou. Model odměny i dvojice preferencí pokrývají jen to, co lidé viděli; mimo tuto oblast nic nebrání tomu, aby se model vydal kamkoli. Trest KL(π_θ ‖ π_ref) ho drží blízko modelu, o kterém víme, že rozumně mluví.

π_ref = mix anotátorů (0,6; 0,25; 0,15) pro správně / „nevím“ / špatně, odměny r = (1; 0; −1) — pořadí z našich preferencí, čísla zvolená pro ukázku. Pro β = 2:

(Barvy podle vzorce pro π* výše.)

Teorie říká totéž co náš pokus: čím menší β, tím dál od výchozího modelu (větší KL) a tím víc odměny. Nesmysly by ale v optimu nepřibývaly: π* je π_ref přenásobené odměnou, a model po SFT říkal nesmysl jen ve 2 % případů — kdyby za něj byla odměna nízká jako za chybu, ještě by ho ubylo. U skutečného DPO přibyl proto, že se učí jen z konečně mnoha dvojic a odměna nesmyslu v nich nikde není: ztráta tlačí dolů horší odpovědi, ale kam se uvolněná pravděpodobnost přesune, neurčuje. Model, který se vzdálí daleko (malé β), pak skončí jinde, než teorie slibuje.

Všech možných odpovědí je příliš mnoho na to, abychom KL sečetli přesně. Použijeme druhý tvar z definice KL, průměr přes odpovědi losované z π_θ. Nechť x₁, …, x_n ∈ Σ* jsou dotazy, n ≥ 1, a pro každé i = 1, …, n je y_i ∼ π_θ(· | x_i) nezávisle vylosovaná odpověď. Odhad KL ze vzorku je číslo

KL̂ = 1/n ∑_i=1ⁿ [ ln π_θ(y_i | x_i) − ln π_ref(y_i | x_i) ].

Jeho střední hodnota je 1/n ∑_i=1ⁿ KL(π_θ(· | x_i) ‖ π_ref(· | x_i)), tedy průměrná KL přes dané dotazy.

u nás n = 240 (6 měst × 40 odpovědí), ln π(y) jako součet přes znaky odpovědi. Odhad je nestranný, ale jednotlivé členy mohou být záporné, a když jsou si modely velmi blízko, může vyjít záporný i průměr: u β = 10 vyšlo −0,031, v tabulce zaokrouhleno na nulu (skutečná KL zápornou být nemůže).