معماری Transfusion: ادغام متن اتورگرسیو و تصویر دیفیوژن در یک ترنسفورمر

مدل‌های چندحالته‌ی نسل قبلی مثل Chameleon ناچار بودند تصاویر را با روش‌های کوانتیزاسیون (مانند VQ-VAE) به توکن‌های گسسته خرد کنند تا ترنسفورمر بتواند با آن‌ها مثل واژه‌های متن رفتار کند. این کار به قیمت افت شفافیت و کیفیت بصری تصویر تمام می‌شد؛ در حالی که مدل‌های اختصاصی تصویر بر پایه‌ی دیفیوژن پیوسته، کیفیتی بسیار بالاتر تولید می‌کردند اما جدا از مدل زبانی بودند.

معماری Transfusion با ترکیب نوآورانه‌ی دو تابع زیان (پیش‌بینی توکن بعدی برای متن و تطبیق جریان/دیفیوژن برای تصویر) روی یک بدنه‌ی ترنسفورمر مشترک، مرز میان پردازش پیوسته و گسسته را حذف می‌کند. در این درس ساختار داخلی این معماری، طراحی ماسک اتنشن مثلثی-بلوکی و پیوند آن با MMDiT در Stable Diffusion 3 را گام‌به‌گام پیاده‌سازی و بررسی می‌کنید.

مدل‌های Chameleon و Emu3 همه‌چیز را روی توکن‌های گسسته شرط بستند. این مدل‌ها کار می‌کنند، اما گلوگاه کوانتیزاسیون (quantization bottleneck) کاملاً مشهود است؛ کیفیت تصویر در سطحی پایین‌تر از مدل‌های دیفیوژن در فضای پیوسته متوقف می‌شود. Transfusion (متا، Zhou و همکاران، آگوست 2024) شرط معکوسی می‌بندد: تصاویر را پیوسته نگه دارید، VQ-VAE را کاملاً کنار بگذارید و یک ترنسفورمر واحد را با دو تابع زیان (loss) آموزش دهید. توکن‌های متنی با تابع پیش‌بینی توکن بعدی (NTP) آموزش می‌بینند و پچ‌های تصویر تابع زیان تطبیق جریان (flow matching) / دیفیوژن دریافت می‌کنند. هر دو هدف، وزن‌های یکسانی را بهینه‌سازی می‌کنند. معماری زیربنایی Stable Diffusion 3 یعنی MMDiT خویشاوند نزدیک همین ایده است. در این درس ایده‌ی اصلی Transfusion را بررسی می‌کنیم، یک ساختار آموزشی اسباب‌بازی با دو تابع زیان می‌سازیم و ماسک اتنشن را که به یک ترنسفورمر اجازه می‌دهد هر دو کار را انجام دهد، گام‌به‌گام ردیابی می‌کنیم.

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

• پیاده‌سازی ترنسفورمری که دو تابع زیان (NTP روی توکن‌های متن و MSE دیفیوژن روی پچ‌های تصویر) را روی یک بک‌بون اجرا می‌کند.
• توضیح این‌که چرا توجه دوطرفه (bidirectional attention) روی پچ‌های تصویر به‌همراه اتنشن علّی (causal attention) روی توکن‌های متن، چینش ماسک درستی است.
• مقایسه‌ی رویکرد Transfusion (تصاویر پیوسته، تابع زیان دیفیوژن) با رویکرد Chameleon (تصاویر گسسته، NTP) از نظر بار پردازشی، کیفیت و پیچیدگی کد.
• توضیح دستاورد MMDiT: وزن‌های اختصاصیِ هر مدالیته در هر بلاک و اتنشن مشترک در جریان پسماند (residual stream).

صورت مسئله

بحث بر سر توکن‌های گسسته در برابر پیوسته‌ی تصویر قدیمی‌تر از مدل‌های زبانی بزرگ است. بازنمایی‌های پیوسته (پیکسل‌های خام، متغیرهای نهفته‌ی VAE) جزئیات را حفظ می‌کنند. توکن‌های گسسته (اندیس‌های VQ) با واژگان ذاتی ترنسفورمر همخوانی دارند، اما در مرحله‌ی کوانتیزاسیون بخشی از جزئیات را هدر می‌دهند.

مدل‌های Chameleon و Emu3 مسیر گسسته را انتخاب کردند: یک تابع زیان، یک معماری، اما با سقفی برای شفافیت تصویر که با کیفیت توکنایزر محدود می‌شد.

مدل‌های دیفیوژن مسیر پیوسته را در پیش گرفتند: کیفیت تصویر فوق‌العاده، اما با مدلی کاملاً مجزا از LLM، مهندسی پیچیده‌ی برنامه‌ی نویز (noise schedule) و نبود ادغام یکپارچه با تولید متن.

پاسخ Transfusion به این چالش: آیا می‌شود هر دو را داشت؟ تصاویر را پیوسته نگه داریم، همچنان یک مدل واحد را آموزش دهیم و دو تابع زیان را در یک گام گرادیان به هم بدوزیم.

مفاهیم پایه

معماری با دو تابع زیان

یک ترنسفورمرِ فقط-دیکودر دنباله‌ای را پردازش می‌کند که شامل این موارد است:

• توکن‌های متن (گسسته، برگرفته از واژگان BPE).
• پچ‌های تصویر (پیوسته، بلاک‌های 16x16 پیکسلی که از طریق یک امبدینگ خطی به بُعد نهفته تصویر می‌شوند — درست مثل ورودی انکودر ViT).
• تگ‌های <image> و </image> که محل قرارگیری پچ‌های پیوسته را مشخص می‌کنند.

گذر رو به جلو یک‌بار اجرا می‌شود. تابع زیان برای هر توکن یکی از دو سر را برمی‌گزیند:

• برای توکن‌های متن: آنتروپی متقاطع استاندارد روی سر لاجیت‌های واژگان.
• برای پچ‌های تصویر: تابع زیان دیفیوژن روی پچ‌های پیوسته — پیش‌بینی نویزی که به هر پچ اضافه شده بود.

گرادیان از بدنه‌ی مشترک ترنسفورمر عبور می‌کند. هر دو تابع زیان هم‌زمان وزن‌های مشترک را بهبود می‌دهند.

ماسک اتنشن: متنِ علّی + تصویرِ دوطرفه

توکن‌های متن باید علّی باشند؛ نمی‌توانید اجازه دهید یک توکن متنی به متن‌های آینده نگاه کند، وگرنه وادارسازی معلم (teacher forcing) به هم می‌ریزد. در مقابل، پچ‌های تصویر نمایانگر یک اسنپ‌شات هستند؛ آن‌ها باید داخل همان بلاک تصویر به‌صورت دوطرفه به یکدیگر توجه کنند.

ماسک:

M[i, j] = 1 if:
  (i is text and j is text and j <= i)   # causal for text
  OR (i is image and j is image and same_image_block(i, j))   # bidirectional within image
  OR (i is text and j is image and j < i)   # text attends to images that closed before it
  OR (i is image and j is text and j < i_image_start)   # image attends to preceding text

این ساختار در فاز آموزش و استنتاج به‌صورت یک ماسک مثلثی-بلوکی پیاده‌سازی می‌شود.

تابع زیان دیفیوژن درون ترنسفورمر

تابع زیان دیفیوژن ساختاری استاندارد دارد: اضافه کردن نویز به یک پچ تصویر و واداشتن مدل به پیش‌بینی نویز (یا مشابهاً پیش‌بینی پچِ تمیز). نسخه‌ی Transfusion از تطبیق جریان استفاده می‌کند — پیش‌بینی میدان سرعت از وضعیت نویزی به تمیز.

در طول آموزش:

1. برای هر پچ تصویر $x_0$، یک گام زمانی تصادفی $t$ نمونه‌گیری کنید.
2. نویز $\epsilon$ را نمونه‌گیری کرده و $x_t = (1-t) \cdot x_0 + t \cdot \epsilon$ را محاسبه کنید (درون‌یابی خطی برای تطبیق جریان).
3. ترنسفورمر مقدار $v_\theta(x_t, t)$ را پیش‌بینی می‌کند؛ $\text{loss} = \mathop{\mathrm{MSE}}(v_\theta(x_t, t), \epsilon - x_0)$.
4. پس‌انتشار هم‌گام با تابع زیان‌های NTP متن از همان دنباله انجام می‌شود.

در فاز استنتاج، فرآیند تولید به این شکل است:

• توکن‌های متن: نمونه‌گیری اتورگرسیو استاندارد.
• پچ‌های تصویر: حلقه‌ی نمونه‌گیری دیفیوژن (معمولاً 10 تا 30 گام) مشروط به توکن‌های متنی پیشین.

معماری MMDiT: روایت Stable Diffusion 3

پروژه‌ی Stable Diffusion 3 (Esser و همکاران، مارس 2024) تقریباً هم‌زمان با Transfusion مدل MMDiT (ترنسفورمر دیفیوژن چندحالته) را عرضه کرد. این معماری‌ها خویشاوند یکدیگرند.

تفاوت‌های کلیدی MMDiT:

• وزن‌های اختصاصی برای هر مدالیته در هر بلاک. هر بلاک ترنسفورمر دارای وزن‌های مجزای Q، K، V و MLP برای توکن‌های متن در برابر پچ‌های تصویر است. اتنشن مشترک است (بین‌مدالیته‌ای)؛ بقیه‌ی بخش‌ها برای هر مدالیته اختصاصی هستند.
• آموزش مبتنی بر جریان تصحیح‌شده. نوع خاصی از تطبیق جریان با رفتار نمونه‌گیری شناخته‌شده و ریاضیات ساده‌تر نسبت به DDPM.
• مقیاس. مدل MMDiT بک‌بون SD3 است (نسخه‌های 2 و 8 میلیارد پارامتری). مقاله‌ی Transfusion تا مقیاس 7B پیش می‌رود.

هر دو روی یک ایده‌ی محوری همگرا می‌شوند: یک ترنسفورمر واحد که NTP را روی متن و دیفیوژن را روی بازنمایی‌های پیوسته‌ی تصویر اجرا می‌کند.

چرا این روش از سبک Chameleon بهتر است؟

شکاف کیفیتی بین دیفیوژن پیوسته و NTP گسسته در تولید تصویر کاملاً ملموس است. مقاله‌ی Transfusion گزارش می‌دهد:

• در مقیاس 7B پارامتر، شاخص FID نسبت به یک مدل هم‌اندازه با معماری Chameleon بین 3 تا 5 امتیاز بهتر است.
• نیازی به آموزش توکنایزر نیست — انکودر تصویر بسیار ساده‌تر است (تصویر کردن خطی به فضای نهفته، درست مثل لایه‌ی ورودی ViT).
• برخلاف توکن‌های اتورگرسیو تصویر، فاز استنتاج می‌تواند عملیات نویززدایی پچ‌های تصویر را موازی‌سازی کند.

نقطه‌ی ضعف: Transfusion مدلی با دو تابع زیان است و همین دینامیک آموزش را پیچیده‌تر می‌کند. وزن‌های توابع زیان نیازمند تنظیم دقیق هستند. ناهماهنگی در زمان‌بندی NTP و دیفیوژن می‌تواند باعث غلبه‌ی یک سر بر دیگری شود.

شاخه‌های مشتق‌شده

مدل Janus-Pro (درس 12.15) ایده‌ی Transfusion را با جدا کردن انکودر بینایی برای درک و تولید پالایش می‌کند — استفاده از SigLIP برای یکی و VQ برای دیگری — در حالی که بدنه‌ی ترنسفورمر مشترک می‌ماند. مدل Show-o (درس 12.14) دیفیوژن را با دیفیوژن گسسته (پیش‌بینی ماسک‌شده) جایگزین می‌کند. خانواده‌ی تولید یکپارچه بعد از Transfusion به‌سرعت منشعب شد.

مدل‌های مالکیتی و بسته‌ی محیط عملیاتی که تصاویر را به‌طور بومی تولید می‌کنند ممکن است از نسخه‌ای از همین خانواده استفاده کنند، اما معماری آن‌ها رسماً منتشر نشده است؛ بنابراین هر ادعایی در این زمینه صرفاً حدس و گمان است.

استفاده از ابزارهای واقعی

فایل code/main.py یک نسخه‌ی اسباب‌بازی از Transfusion را روی مسئله‌ای کوچک و شبیه MNIST پیاده‌سازی می‌کند:

• متن‌های توصیفی (captions) دنباله‌های کوتاهی از اعداد صحیح هستند که یک رقم (0 تا 9) را توصیف می‌کنند.
• تصاویر شبکه‌های 4x4 از بایت‌ها هستند.
• یک جفت تصویر کردن خطی با وزن مشترک نقش جانشین ترنسفورمر را ایفا می‌کند؛ تابع زیان NTP روی متن، و تابع زیان MSE روی پچ‌های نویزی.
• حلقه‌ی آموزش بین دو تابع زیان تناوب برقرار می‌کند و ماسک اتنشن به‌طور صریح تعریف می‌شود.
• یک پیش‌نمایش از ماسک، الگوی اتنشن مثلثی-بلوکی را برای دنباله‌ای از نوع «متن، تصویر، متن» چاپ می‌کند. مرحله‌ی تولید اجرا نمی‌شود: یک نمونه‌گیر واقعی متن را توکن‌به‌توکن دیکود می‌کند و سپس پچ‌های تصویر را در گام‌های متعدد نویززدایی می‌کند، درست همان‌طور که در تمرین‌های زیر محاسبه می‌کنید.

این ترنسفورمر در حد یک اسباب‌بازی است. دستاورد واقعی این کد، ساختار لوله‌کشیِ دو تابع زیان، ساخت ماسک اتنشن و طراحی حلقه‌ی استنتاج است.

خروجی این درس

این درس سند outputs/skill-two-loss-trainer-designer.md را تولید می‌کند. این خروجی با دریافت یک تسک آموزشی چندحالته‌ی جدید (متن + تصویر، متن + صوت، متن + ویدیو)، برنامه‌ی مدیریت دو تابع زیان (وزن‌های زیان، شکل ماسک، بلاک‌های مشترک در برابر اختصاصیِ مدالیته) را طراحی کرده و ریسک‌های پیاده‌سازی را مشخص می‌کند.

تمرین‌ها

1. مدلی به سبک Transfusion روی 70% توکن‌های متنی و 30% پچ‌های تصویری آموزش می‌بیند. اندازه‌ی عددی تابع زیان دیفیوژن تصویر تقریباً 10 برابر تابع زیان NTP متن است. چه وزن‌هایی برای توابع زیان، تعادل را برقرار می‌کنند؟

2. ماسک مثلثی-بلوکی را برای دنباله‌ی [T, T, <image>, P, P, P, P, </image>, T] پیاده‌سازی کنید. هر درایه را با 0 یا 1 مشخص کنید.

3. مدل MMDiT از وزن‌های QKV مجزا برای هر مدالیته استفاده می‌کند. این ساختار در مقایسه با ترنسفورمر کاملاً مشترک Transfusion چه میزان بار اضافه به تعداد پارامترها تحمیل می‌کند؟ در مقیاس 7B، آیا این هزینه ارزشش را دارد؟

4. در فاز تولید: با داشتن یک پرامپت متنی، مدل ابتدا NTP را به مدت 50 توکن اجرا می‌کند، سپس به <image> می‌رسد و دیفیوژن را روی 256 پچ در طول 20 گام نویززدایی اجرا می‌کند. در مجموع چند گذر رو به جلو لازم است؟

5. بخش 3 مقاله‌ی SD3 را بخوانید. مفهوم جریان تصحیح‌شده (rectified flow) را شرح دهید و بگویید چرا در گام‌های استنتاج کمتری نسبت به DDPM همگرا می‌شود.

اصطلاحات کلیدی

برای مطالعه‌ی بیشتر

• Zhou et al. — Transfusion (arXiv:2408.11039)
• Esser et al. — Stable Diffusion 3 / MMDiT (arXiv:2403.03206)
• Peebles & Xie — DiT (arXiv:2212.09748)
• Zhao et al. — MonoFormer (arXiv:2409.16280)
• Xie et al. — Show-o (arXiv:2408.12528)