‏Lighthouse Attention: آموزش مدل‌های زبانی با کانتکست ۵۱۲ هزار توکنی

بزرگ‌ترین مشکل آموزش مدل‌های زبانی با کانتکست طولانی اینه که attention استاندارد هزینه محاسباتی درجه دوم (quadratic) داره — یعنی با دو برابر شدن طول کانتکست، هزینه چهار برابر می‌شه. FlashAttention این هزینه رو کمی کاهش داده، ولی اساس مشکل همچنان پابرجاست. Lighthouse Attention رویکرد متفاوتی داره: به جای بهینه‌سازی کرنل attention، یه زیردنباله کوچیک و متراکم انتخاب می‌کنه و FlashAttention معمولی رو روی اون اجرا می‌کنه.

ایده اصلی اینه که Q، K و V همه با هم به صورت متقارن در یه هرم چندسطحی (multi-resolution pyramid) ادغام (pool) می‌شن. در هر سطح، توکن‌های مجاور با هم ترکیب می‌شن تا خلاصه‌ای از بازه‌های بزرگ‌تر بسازن. بعد، با استفاده از نرم ℓ₂ هر بردار (بدون هیچ پارامتر قابل یادگیری اضافه‌ای)، مهم‌ترین ورودی‌های هرم انتخاب می‌شن و FlashAttention فقط روی این زیرمجموعه اجرا می‌شه.

یکی از نکات کلیدی طراحی اینه که انتخاب خارج از کرنل attention انجام می‌شه. نتیجه‌ی انتخاب یه tensor متراکم و پیوسته‌ست، نه یه ساختار sparse با ایندکس‌های پراکنده. این یعنی می‌شه بدون هیچ تغییری در کرنل FlashAttention ازش استفاده کرد. اکثر روش‌های مشابه مثل NSA یا MoBA مجبورن کرنل sparse مخصوص خودشون رو بنویسن — Lighthouse این نیاز رو حذف می‌کنه.

آموزش دو مرحله داره: در مرحله اول مدل با Lighthouse Attention آموزش می‌بینه (بیشتر بودجه). در مرحله دوم همون checkpoint با attention معمولی (SDPA) ادامه می‌ده تا مدل به یه مدل dense کامل تبدیل بشه. در هر سه نقطه‌ی مختلف از این resume، مدل نهایی به اندازه مدل dense-from-scratch یا بهتر از اون بود — این نشون می‌ده آموزش sparse توانایی dense attention مدل رو خراب نمی‌کنه.

نتایج روی مدل ۵۳۰M Llama-3 با ۱۶ هزار گام بهینه‌سازی و حدود ۵۰ میلیارد توکن تایید شده. سرعت آموزش در مرحله اول بین ۸۴ تا ۱۲۶ هزار توکن در ثانیه در هر GPU بوده، در مقابل ۴۶ هزار برای attention معمولی. برای کانتکست‌های بسیار بلند هم با Context Parallelism روی ۳۲ GPU از سری Blackwell، آموزش با ۱ میلیون توکن کانتکست موفق بوده.

نکات کلیدی:

• Lighthouse Attention در کانتکست ۵۱۲K توکنی ~۱۷ برابر سریع‌تر از attention معمولی روی یک B200 اجرا می‌شه
• Q، K و V همه به صورت متقارن pool می‌شن — برخلاف روش‌های مشابه که فقط K و V رو pool می‌کنن
• انتخاب با نرم ℓ₲ پارامتر اضافه، scorer قابل یادگیری، یا Gumbel-softmax نداره
• خروجی انتخاب یه tensor متراکم معمولیه — FlashAttention بدون تغییر روش کار می‌کنه
• پس از دو مرحله آموزش، مدل نهایی یه dense attention model کامله و در inference نیازی به Lighthouse نداره
• کد در گیت‌هاب منتشر شده: github.com/ighoshsubho/lighthouse-attention