‏Recurrent Transformer: عمق بیشتر، دیکود سریع‌تر

ترنسفورمرها قلب تپنده‌ی اکثر مدل‌های زبانی امروزی‌ان، ولی یه مشکل اساسی دارن: اون‌ها توکن‌ها رو به‌صورت موازی پردازش می‌کنن و از نظر «عمق زمانی» محدودیت دارن. به زبان ساده، در هر موقعیت از دنباله، هر لایه فقط می‌تونه به خروجی لایه‌ی قبلی نگاه کنه، نه به تاریخچه‌ی کامل پردازش. این یعنی عمق مؤثر مدل در طول زمان، گره خورده به تعداد لایه‌هاست و نمی‌شه فراتر رفت.

مدل‌های بازگشتی (Recurrent) مثل RNN ها این محدودیت رو ندارن و می‌تونن عمق زمانی نامحدودی داشته باشن، ولی مشکلات خودشون رو دارن: بهینه‌سازی ناپایدار و ضعف در بهره‌گیری از شتاب‌دهنده‌های مدرن مثل GPU ها باعث شده بودن که از رقابت با ترنسفورمرها عقب بمونن.

حالا یه تیم تحقیقاتی یه معماری ساده و جذاب معرفی کرده به اسم Recurrent Transformer. ایده‌ی اصلی اینه که هر لایه به جای اینکه فقط به خروجی لایه‌ی قبلی توجه کنه، به key-value هایی که از فعال‌سازی‌های خودش تولید شده هم توجه داشته باشه. این تغییر کوچیک یه نوع حافظه‌ی بازگشتی لایه‌به‌لایه ایجاد می‌کنه بدون اینکه هزینه‌ی دیکود autoregressive استاندارد زیاد بشه.

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

یه چالش عملی مهم در این معماری اینه که در مرحله‌ی prefill و آموزش، چون کلیدها و مقادیر به‌صورت ترتیبی آشکار می‌شن، ظاهراً فرآیند bandwidth-bound می‌شه و arithmetic intensity مؤثر نزدیک به ۱ می‌مونه. ولی محققان یه الگوریتم دقیق مبتنی بر tiling ارائه دادن که ترافیک HBM (حافظه‌ی پرسرعت GPU) رو از مرتبه‌ی N² به N log N کاهش می‌ده و arithmetic intensity رو به مرتبه‌ی N/log N می‌رسونه. این بهبود کاملاً ریاضیاتی و بدون تغییر در محاسبات اصلیه.

آزمایش‌ها روی مدل‌هایی با ۱۵۰ و ۳۰۰ میلیون پارامتر روی مجموعه داده‌ی C4 انجام شده. نتایج نشون می‌ده Recurrent Transformer نسبت به ترنسفورمر معمولی با همون تعداد پارامتر، cross-entropy بهتری داره و جالب‌تر اینکه این بهبود رو با لایه‌های کمتری به دست می‌یاره.

این موضوع نشون می‌ده که بازگشت (recurrence) می‌تونه جای عمق رو بگیره و در عوض مدل عریض‌تر بشه. این trade-off یه مزیت عملی مهم داره: حجم KV cache کمتر می‌شه که مستقیماً به کاهش مصرف حافظه در زمان inference و کاهش تأخیر تولید توکن منجر می‌شه.

نکات کلیدی:

• ترنسفورمرهای معمولی از نظر عمق زمانی محدودیت دارن و Recurrent Transformer این مشکل رو با اضافه کردن حافظه‌ی بازگشتی لایه‌ای حل می‌کنه
• معماری جدید می‌تونه هم رفتار ترنسفورمر و هم رفتار مدل‌های بازگشتی رو شبیه‌سازی کنه
• یه الگوریتم tiling-based اختصاصی، پیچیدگی ترافیک حافظه رو از N² به N log N کاهش می‌ده
• نتایج روی C4 نشون می‌ده این معماری با لایه‌های کمتر عملکرد بهتری داره
• KV cache کوچک‌تر یعنی مصرف حافظه‌ی کمتر و inference سریع‌تر در دنیای واقعی
• این رویکرد از بی‌ثباتی بهینه‌سازی که گریبان مدل‌های بازگشتی قدیمی رو می‌گرفت فرار کرده