🔬 Метод

Первое явление, которое обнаруживают авторы - это смещение внимания от sink токенам к последним по ходу драфтинга (attention drift). Причем оно имеет место как и для EAGLE-3, так и для MTP. Кроме того, sink-like поведение и похожий сдвиг есть даже у моделей, специально заточенных бороться с sink-ами (Qwen3-Next, GPT-OSS).

Чтобы разобраться в причинах возникновения, авторы замеряют нормы входных состояний в драфтер, скрытого состояния и выходов. Обнаруживают, что состояния из более поздних слоев имеют большую норму и могут на себя оттягивать внимание драфтера. Кроме того, нормы состояния драфтера и выходов растут 🔼 с ростом глубины спекуляции. Последний эффект, вероятно, и вызывает перекос attention, ибо более высокие нормы ключей/значений перетягивают на себя вероятностную массу.

Отсюда возникает идея 🤷‍♂️ заменить pre-norm на post-norm в драфтере, чтобы поддерживать норму состояния примерно постоянной. Данная модификация заметно улучшает принятие с увеличением глубины драфта.

Кроме того пробуют добавить gated attention, чтобы подавить влияние синка. Синк оно подавляет и attention drift, но нормы продолжают расти.

Затем еще показывают, что post норма делает выход гораздо более устойчивым к шуму.

🧪 Эксперименты

Итоговая модификация выглядит следующим образом:
📌 Добавляем RMS норм на входные хиддены таргета
📌 Заменяем pre-norm на post-norm для выходного состояния драфта

Переход на post-norm повышает acceptance rate на 14% в среднем по моделям в ID режиме. Шатания шаблона заметно просаживают принятие у pre-norm, но почти не влияют на post-norm.

Gated attention делает модель более чувствительной к системному промпту, поэтому в итоге не используется.

Далее пробуют решить проблему с обобщаемостью на длинный контекст. Дефолтный EAGLE-3 ломается, acceptance уходит в 0. С post-norm уже заметно лучше, но все еще менее 1 принятого токена. А вот что реально позволяет обобщаться - это переход от полного attention к SWA (внимание только на скользящее окно последних токенов) + BOS (или системный промпт). Таким образом удается поддерживать acceptance length около 3.

💡 Выводы

Как будто несложная, но хорошо мотивированная, подкрепленная анализом модификация EAGLE. Более того, это уже даже выкатили в vLLM под названием EAGLE-3.1 (см. блогпост) и обучили для Kimi K 2.6. Интересно было бы увидеть честное сравнение с DFlash.