آموزش پیاده‌سازی DPO از صفر: بهینه‌سازی مستقیم ترجیحات در پایتون

پس از مرحله‌ی تنظیم دقیق نظارت‌شده (SFT)، مدل‌های زبانی معمولاً دستورالعمل‌ها را درک می‌کنند، اما لحن و کیفیت پاسخ‌هایشان نامتوازن باقی می‌ماند. روش سنتی برای هم‌راستاسازی خروجی این مدل‌ها با ترجیحات انسانی، خط لوله‌ی پیچیده و پرهزینه‌ی RLHF مبتنی بر الگوریتم PPO است که نگهداری و آموزش هم‌زمان چند مدل (خط‌مشی، پاداش و مرجع) را روی کارت گرافیک می‌طلبد.

الگوریتم بهینه‌سازی مستقیم ترجیحات (Direct Preference Optimization یا DPO) با اثبات ریاضی نشان می‌دهد که نیازی به مدل پاداش مجزا نیست. DPO مسأله‌ی هم‌راستاسازی را به یک تابع زیان نظارت‌شده تبدیل می‌کند که مستقیماً روی جفت‌داده‌های ترجیحی محاسبه می‌شود و پایداری فوق‌العاده‌ای در آموزش به همراه دارد.

در این درس، پایه‌های ریاضی فرمول DPO را استخراج می‌کنیم و یک حلقه‌ی کامل آموزش ترجیحات را با یک مدل ساده در پایتون می‌سازیم. این پیاده‌سازی پایه به‌گونه‌ای طراحی شده که حتی روی CPU یا سخت‌افزارهای معمولی شخصی به‌راحتی اجرا شود و بدون نیاز به منابع پردازشی سنگین، شیوه‌ی کار DPO را نشان دهد.

مدل‌های پاداش و PPO پایپ‌لاین کلاسیک یادگیری تقویتی بر اساس بازخورد انسانی (RLHF) هستند. DPO این ساختار را در یک تابع زیان (loss) نظارت‌شده‌ی واحد خلاصه می‌کند که خط‌مشی (policy) را مستقیماً بر اساس جفت ترجیحات (preference pairs) آموزش می‌دهد. در این درس، تابع زیان DPO را بر مبنای اتحاد اختلاف پاداش اثبات ریاضی (derivation) می‌کنیم، یک مدل مرجع (reference model) و مدل خط‌مشیِ آماده‌به‌کار می‌سازیم، لگاریتم احتمالات را به‌ازای هر توکن به دست می‌آوریم و یک ترنسفورمر بسیار کوچک را روی داده‌های آماده‌ی ترجیحات، شامل پاسخ‌های منتخب (chosen) و مردود (rejected)، آموزش می‌دهیم. آزمون‌ها صحت محاسبات ریاضیِ زیان و جهت گرادیان را تثبیت می‌کنند تا مطمئن شوید پیاده‌سازی‌تان دقیقاً با مقاله مطابقت دارد.

اهداف یادگیری

• اثبات ریاضی تابع زیان DPO به‌صورت تابعی سیگموئید روی اختلاف نسبتِ لگاریتمیِ مقیاس‌بندی‌شده، و ربط دادن آن به پاداش ضمنی (implicit reward).
• ساخت زوجِ مدل مرجع + مدل خط‌مشی، با مدل مرجع فریز‌شده و خط‌مشی قابل آموزش.
• محاسبه‌ی لگاریتم احتمالات در سطح توالی برای هر دو مدل، با ماسک کردن توکن‌های پرامپت.
• آموزش خط‌مشی روی سه‌گانه‌های (prompt, chosen, rejected) و مشاهده‌ی افزایش لگاریتم احتمال پاسخ منتخب نسبت به پاسخ مردود.
• تثبیت رفتار پیاده‌سازی با نوشتن آزمون برای محاسبات ریاضی زیان، علامت گرادیان و تغییرناپذیری مدل مرجع.

صورت مسئله

شما یک مدل تنظیم دقیق نظارت‌شده (SFT) دارید. این مدل دستورالعمل‌ها را دنبال می‌کند، اما خروجی‌های یکدستی ندارد؛ بعضی پاسخ‌ها دقیق و روشن‌اند و بعضی دیگر پرحرف، بی‌ربط یا اشتباه. در عین حال مجموعه‌داده‌ی کوچکی از جفت ترجیحات در اختیار دارید: برای هر پرامپت، یک انسان پاسخ بهتر را به‌عنوان منتخب و دیگری را به‌عنوان مردود علامت‌گذاری کرده است.

پاسخ کلاسیک RLHF یک پایپ‌لاین دو‌مرحله‌ای است: اول یک مدل پاداش (reward model) روی ترجیحات آموزش می‌دهید، بعد با استفاده از PPO، خط‌مشی را نسبت به این پاداش بهینه‌سازی می‌کنید. این کار شدنی است اما هزینه دارد: دو مدل هم‌زمان در حافظه هنگام اجرای PPO، کنترل جریمه‌ی KL برای حفظ خط‌مشی در نزدیکی مدل مرجع، و خطر دور زدن پاداش (reward hacking) در صورت شکننده بودن مدل پاداش.

روش DPO هر دو مرحله را با یک تابع زیان نظارت‌شده‌ی واحد جایگزین می‌کند. مدل پاداش اصلاً وجود خارجی مستقل پیدا نمی‌کند. خط‌مشی مستقیماً روی جفت ترجیحات آموزش می‌بیند و جریمه‌ی KL نسبت به مرجعِ SFT هم مستقیماً در فرمول لحاظ شده است. نتیجه تحت مدل ترجیحات بردلی-تری (Bradley-Terry) همان راه‌حل بهینه است، اما با کدی به‌مراتب کمتر.

مبانی نظری

نقطه‌ی شروع، مدل بردلی-تری است. به‌ازای پرامپت x و دو پاسخِ y_w (منتخب) و y_l (مردود)، احتمال این‌که انسان y_w را ترجیح دهد برابر است با:

$$P(y_w > y_l \mid x) = \sigma\big( r(x, y_w) - r(x, y_l) \big)$$

که در آن r تابع پاداش پنهان است. در RLHF ابتدا r را بر اساس ترجیحات برازش می‌دهند و بعد خط‌مشی pi را طوری آموزش می‌دهند که r را همراه با مهارکننده‌ی KL بیشینه کند:

$$\max_{\pi} \; \mathbb{E}_{x, y \sim \pi} \big[ r(x, y) \big] - \beta \, \mathrm{KL}(\pi \,\|\, \pi_{\text{ref}})$$

در اثبات ریاضی DPO می‌بینیم که خط‌مشی بهینه‌ی *pi تحت این هدف، یک فرم بسته (closed form) بر حسب r دارد:

$$\pi^*(y \mid x) = \frac{1}{Z(x)} \, \pi_{\text{ref}}(y \mid x) \exp\!\left( \frac{r(x, y)}{\beta} \right)$$

رابطه را بر حسب r بازنویسی می‌کنیم:

$$r(x, y) = \beta \big( \log \pi^*(y \mid x) - \log \pi_{\text{ref}}(y \mid x) \big) + \beta \log Z(x)$$

عبارت log Z(x) برای هر دو پاسخ y_w و y_l یکسان است (چون فقط به x وابسته است، نه y)؛ بنابراین هنگام محاسبه‌ی تفاضل ترجیحات خط می‌خورد:

$$\begin{aligned} r(x, y_w) - r(x, y_l) = \beta \big( &\log \pi_\theta(y_w \mid x) - \log \pi_{\text{ref}}(y_w \mid x) \\ &- \log \pi_\theta(y_l \mid x) + \log \pi_{\text{ref}}(y_l \mid x) \big) \end{aligned}$$

این عبارت را درون تابع سیگموئیدِ بردلی-تری می‌گذاریم و منفی لگاریتم درست‌نمایی (negative log likelihood) را روی جفت ترجیحات می‌گیریم:

$$\begin{aligned} L_{\text{DPO}}(\theta) = -\mathbb{E}_{(x, y_w, y_l)} \Big[ \log \sigma\big( \beta \big( &\log \pi_\theta(y_w \mid x) - \log \pi_{\text{ref}}(y_w \mid x) \\ &- \log \pi_\theta(y_l \mid x) + \log \pi_{\text{ref}}(y_l \mid x) \big) \big) \Big] \end{aligned}$$

این همان تابع زیان است: یک سیگموئید روی یک مقدار اسکالرِ منفرد به‌ازای هر نمونه، که از چهار لگاریتم احتمال به دست می‌آید. نه مدل پاداش جداگانه‌ای در کار است، نه PPO، و نه حتی عبارت KL جداگانه در تابع زیان؛ چون قید KL مستقیماً درون اثباتِ فرم بسته تعبیه شده است.

علامت گرادیان

پیش از شروع هر دور آموزش، یک بررسی منطقی (sanity check) حسابی به کار می‌آید. گرادیان را نسبت به log pi_theta(y_w | x) حساب می‌کنیم:

$$\frac{\partial L_{\text{DPO}}}{\partial \log \pi_\theta(y_w \mid x)} = -\beta \big(1 - \sigma(z)\big)$$

که در آن z ورودی تابع سیگموئید است. این مقدار به‌ازای تمام مقادیر z منفی است؛ یعنی بالا بردن لگاریتم احتمال خط‌مشی برای پاسخ منتخب، تابع زیان را کم می‌کند. در جهت قرینه، گرادیان نسبت به log pi_theta(y_l | x) مثبت است: بالا رفتن لگاریتم احتمال پاسخ مردود، تابع زیان را زیاد می‌کند. آموزش، منتخب را بالا می‌کشد و مردود را پایین می‌برد. مدل مرجع هم دست‌نخورده و ثابت می‌ماند.

داده‌ها

دوازده سه‌گانه‌ی ترجیحات همراه این درس آمده است. هر سه‌گانه به‌شکل (prompt, chosen, rejected) تعریف می‌شود. پاسخ منتخب کوتاه و دقیق است و پاسخ مردود پرحرف، بی‌ربط یا غلط. این جفت‌ها همان خانواده‌های وظیفه‌ی درس 39 (پایتخت‌ها، محاسبات ریاضی و لیست‌ها) را پوشش می‌دهند تا خط‌مشی که از پایه‌ی SFT شروع کرده، نقطه‌ی شروع معقولی داشته باشد.

این مجموعه‌داده عمداً کوچک انتخاب شده است. DPO در محیط عملیاتی (پروداکشن) با ده‌ها هزار جفت داده کار می‌کند؛ اما اینجا هدف این است که فرمول‌های ریاضیِ زیان و حلقه‌ی آموزش به‌صورت سرتاسری روی داده‌های بسیار جمع‌وجور اجرا شوند و ببینید که فاصله‌ی لگاریتم احتمال پاسخ منتخب نسبت به مردود چطور شفاف رشد می‌کند.

تغییرناپذیری مدل مرجع

در پیاده‌سازی DPO باید مدل مرجع را با وسواس مدیریت کرد. مدل مرجع در واقع همان مدل SFT است که فریز شده و تکان نمی‌خورد. سه ویژگی باید برقرار باشد:

• پارامترهای مدل مرجع هرگز گرادیان دریافت نکنند.
• لگاریتم احتمالات مدل مرجع در طول دوره‌ها (epochs) هیچ تغییری نکند.
• وزن‌های اولیه‌ی خط‌مشی دقیقاً همان وزن‌های مدل مرجع باشد. (مقدار بهینه‌ی theta برابر با مدل مرجع به‌اضافه‌ی تغییراتِ یادگرفته‌شده است؛ بنابراین مقداردهی اولیه‌ی خط‌مشی با کپی کردن مرجع، دقیق‌ترین نقطه‌ی شروع است.)

پیاده‌سازی این شروط را این‌طور تضمین می‌کند:

• قرار دادن مدل مرجع در بلوک torch.no_grad() در گذرهای رو به جلو (forward passes).
• تنظیم requires_grad=False روی تک‌تک پارامترهای مدل مرجع.
• ساخت خط‌مشی با policy.load_state_dict(reference.state_dict()) بعد از ایجاد مدل مرجع.

معماری

مدل دقیقاً همان TinyGPT درس 39 است (ترنسفورمر فقط-دیکودر، علّی، با توکن‌ساز بایتی). مدل مرجع و خط‌مشی معماری یکسانی دارند؛ در طول آموزش، وزن‌های خط‌مشی رفته‌رفته از مدل مرجع فاصله می‌گیرند، در حالی که مدل مرجع دست‌نخورده باقی می‌ماند.

چه چیزی خواهید ساخت

این پیاده‌سازی شامل یک فایل main.py به همراه آزمون‌هاست:

1. InstructionTokenizer: توکن‌ساز بایتی با توکن‌های اختصاصی INST و RESP، درست مثل درس 39.
2. TinyGPT: ترنسفورمر فقط-دیکودر با همان ساختار درس 39، تا اگر آن درس را نخوانده باشید هم همه‌چیز کامل و مستقل باشد.
3. make_preferences: دوازده سه‌گانه‌ی (prompt, chosen, rejected) را تولید می‌کند.
4. sequence_log_prob: با گرفتن مدل، پرامپت و پاسخ، مجموع لگاریتم احتمالاتِ توکن بعدی را فقط روی بخش پاسخ حساب می‌کند (بدون دخالت توکن‌های پرامپت).
5. dpo_loss: با گرفتن چهار مقدار لگاریتم احتمال و پارامتر beta، تنسورِ زیان به‌ازای هر نمونه و تفاضل پاداش ضمنی را برای لاگ گرفتن تحویل می‌دهد.
6. train_dpo: حلقه‌ی آموزش در هر دوره که لگاریتم احتمالات پاسخ منتخب و مردود را در خط‌مشی و مدل مرجع حساب می‌کند، تابع زیان را اعمال می‌نماید و یک گام بهینه‌ساز Adam می‌زند.
7. evaluate_margins: میانگین حاشیه‌ی تفاضل لگاریتم احتمال منتخب و مردود را در هر لحظه روی خط‌مشی اندازه می‌گیرد.
8. run_demo: مدل مرجع و خط‌مشی را از یک پیش‌آموزش سبک اولیه می‌سازد، وزن‌ها را کپی می‌کند، 30 گام آموزش می‌دهد، تابع زیان و مقدار حاشیه را در هر گام چاپ می‌کند و در صورت اجرای موفق، با خروجی صفر خارج می‌شود.

چرا DPO جواب می‌دهد؟

روش DPO تا جایگاه پارامترسازی پاداش، از نظر ریاضی دقیقاً با RLHF تحت مدل ترجیحات بردلی-تری معادل است. پاداش ضمنی r(x, y) = beta * (log pi(y|x) - log pi_ref(y|x)) را می‌توان از روی ترجیحات تا حد تابعی از x تخمین زد که آن هم در تفاضل نهایی حذف می‌شود. این فرم بسته‌ی خط‌مشی به شما امکان می‌دهد مدل پاداش صریح را کاملاً کنار بگذارید. قید KL هم به‌طور ساختاری اعمال می‌شود: هر انحراف pi از pi_ref نسبت لگاریتمی را بزرگ‌تر می‌کند و تابع سیگموئید اشباع می‌شود؛ در نتیجه وقتی خط‌مشی خیلی دور شود، گرادیان افت می‌کند و متوقف می‌شود. مدل مرجع، شبکه‌ی نجات شماست.

چالش‌های پیشنهادی

• نرمال‌سازی طول را به مجموع لگاریتم احتمالات اضافه کنید: تقسیم بر طول پاسخ. سوگیری طول (Length bias) یکی از ایرادهای رایج DPO است که در آن مدل ترجیح می‌دهد پاسخ‌های کوتاه‌تر را انتخاب کند، چون لگاریتم احتمال آن‌ها از نظر قدرمطلق بزرگ‌تر است.
• نسخه‌ی IPO از این تابع زیان را اضافه کنید: جای سیگموئید و لگاریتم، عبارت (z - 1)^2 را بگذارید و سرعت همگرایی را روی داده‌های این درس مقایسه کنید.
• پارامتر هموارسازی برچسب (label smoothing) را پیاده کنید تا برچسب‌های صفر و یک منتخب/مردود با عدد یکنواخت 0.5 ترکیب شوند.
• مدل مرجع را با یک مدل کوچک‌تر و سبک‌تر جایگزین کنید (رویکردی با سبکِ تقطیر دانش).

با این پیاده‌سازی، کد تابع زیان، تغییرناپذیری مرجع و حلقه‌ی آموزش را در دست دارید. اصل این درس، ریاضیات آن است و کد، این ریاضیات را برایتان ملموس و عینی می‌کند.