‏Token-In Token-Out؛ تله‌ی پنهان آموزش RL ایجنت‌ها

این مقاله یه باگِ پنهان و موذی در آموزش RL مدل‌های زبانی رو کالبدشکافی می‌کنه. در حالت تک‌مرحله‌ای همه‌چیز عالی به‌نظر می‌رسه: منحنی‌ها تمیز، پاداش‌ها منطقی، همه‌چیز همگرا. ولی وقتی می‌خوای یه ایجنت آموزش بدی و به مدل اجازه می‌دی وسط rollout ابزار صدا بزنه، منحنی‌ها عجیب می‌شن، loss بی‌دلیل جهش می‌کنه و آخرش با یه خطای shape mismatch می‌شکنه.

چیزی که تقریباً قطعاً داره اتفاق می‌افته اینه که حلقه‌ی rollout تو به‌صورت خاموش قاعده‌ی Token-In Token-Out (TITO) رو نقض می‌کنه. تو پاسخ مدل رو پارس می‌کنی تا فراخوانی ابزار رو تشخیص بدی، بعد مکالمه‌ی به‌روزشده رو دوباره توکنایز می‌کنی. معمولاً این رفت‌وبرگشت همون توکن‌های قبلی رو برمی‌گردونه، ولی گاهی نه؛ و اون موقع گرادیان روی دنباله‌ای می‌افته که مدل اصلاً نمونه‌برداری نکرده. کد کرش نمی‌کنه، ولی ریاضیاتش بی‌سروصدا خراب می‌شه و سیگنال گرادیان کاملاً غیرقابل‌اعتماد می‌شه.

اصل ماجرا اینه که در RL باید دقیقاً روی همون توکن‌هایی گرادیان بزنی که مدل تولید کرده — نه پاسخ ابزار، چون اون از policy نیومده. ولی decode عکسِ encode نیست: چند دنباله‌ی توکنیِ متفاوت می‌تونن به یه رشته‌ی یکسان decode بشن، پس اگه توکن‌ها رو decode کنی، متن رو دوباره encode کنی، ممکنه به دنباله‌ی متفاوتی برسی. دلیلش هم اینه که mergeهای BPE روی مرز توکن‌ها پایدار نیستن، به‌علاوه‌ی آزادی‌هایی مثل whitespace در JSON، ترتیب آرگومان‌ها و نحوه‌ی رندر دوباره‌ی توکن‌های ویژه. حلقه‌ای که همه اول می‌نویسن — نگه‌داشتن مکالمه به‌شکل لیست پیام، رندر دوباره در هر نوبت و توکنایزِ کل در آخر — دقیقاً همین‌جا می‌شکنه.

دو راه‌حل هست. اول، یه رابطِ render مخصوص هر مدل بسازی که می‌دونه چطور پیام‌ها رو فرمت و پاسخ‌ها رو پارس کنه بدون رندر دوباره؛ کتابخونه‌ی renderers همین کارو می‌کنه و خانواده‌های اصلی مدل‌های open-weights رو پوشش می‌ده، ولی هزینه‌اش ساختاریه: هر مدل جدید یه renderer دستی می‌خواد. راه‌حل دوم و تمیزتر، یه قاعده‌ی واحده: هیچ‌وقت توکن‌هایی که decode کردی رو دوباره encode نکن. توکن‌های نمونه‌برداری‌شده‌ی مدل مستقیم می‌رن توی یه بافرِ در حال رشد و همون بافر منبع حقیقته؛ لیست پیام‌ها فقط دفترداری می‌شه. آره، توکن‌ها رو پارس می‌کنیم تا بفهمیم باید ابزار صدا بزنیم یا نه، ولی نتیجه‌ی پارس فقط برای مسیریابیه و هیچ‌وقت به پرامپت برنمی‌گرده. با این کار هم مرزهای هر نوبت حفظ می‌شن (چون از اول معلوم بودن) و هم drift توکن‌ها از بین می‌ره. تنها عملیات chat-template باقی‌مونده در حلقه، توکنایزکردن پاسخ ابزار و اضافه‌کردنشه که با تفاضل دو رندرِ ساختگی به‌دست میاد.

نکات کلیدی:

• در آموزش RL ایجنت‌های ابزاردار، حلقه‌ی rollout می‌تونه خاموش قاعده‌ی Token-In Token-Out رو نقض کنه
• علت: tokenization برگشت‌پذیر نیست؛ decode و دوباره encode می‌تونه به توکن‌های متفاوت برسه
• نتیجه: گرادیان روی دنباله‌ای می‌افته که مدل تولید نکرده و سیگنال خراب می‌شه
• راه‌حل اصلی: هیچ‌وقت توکن‌های decode‌شده رو دوباره encode نکن و توکن‌های نمونه‌برداری‌شده رو در یه بافر نگه دار