DFlash ускоряет спекулятивный декодинг с помощью диффузионного драфтера

DFlash — диффузионный метод спекулятивного декодинга Сегодня поговорим о DFlash — методе спекулятивного декодинга, использующем легковесную модель блочной диффузии для параллельной генерации драфтовых токенов. Но для начала о том, почему имеет смысл использовать диффузионные модели в спекулятивном декодинге. Средняя задержка генерации одного токена определяется суммой времени верификации и времени драфтинга, делённой на среднюю длину принятой цепочки, включая бонусный токен. Соответственно, ускорение в спекулятивном декодинге зависит не только от количества принятых токенов, но и от стоимости построения гипотезы. Главная проблема классического авторегрессионного спекулятивного декодинга в том, что время драфтинга линейно растёт с длиной генерируемой цепочки. Поэтому драфтеры обычно делают относительно небольшими, чтобы построение гипотезы занимало разумное время. Диффузионные подходы позволяют предсказывать блок токенов параллельно, за один forward pass, благодаря чему время драфтинга почти не скейлится от размера блока. Именно это — основная мотивация для применения диффузионных моделей в спекулятивном декодинге. С точки зрения архитектуры DFlash — это пятислойный трансформер, каждый слой которого работает с представлениями, полученными из основной модели, и с представлениями расшумляемого блока токенов. На входе первый слой получает последний декодированный токен основной модели, к которому добавляются маскированные токены, предназначенные для предсказания за один forward pass. Контекст основной модели формируется способом, похожим на предложенный в EAGLE-3: берутся скрытые представления с нескольких слоёв основной модели, конкатенируются, после чего к ним применяются линейная проекция и нормализация. Однако здесь, в отличие от EAGLE-3, фичи контекста, подаваемые в каждый из пяти слоёв, одинаковы — меняются только представления расшумляемого блока. Авторы называют это KV injection. Внутри спекулируемого блока используется полный attention, а по отношению к предыдущему контексту внимание остается каузальным. Теперь перейдём к обучению. Авторы не разделяют обучающий пример равномерно на блоки фиксированной длины, а семплируют якорные токены. Каждый из них рассматривается как последний уже декодированный токен, после которого добавляется необходимое количество маскированных. Ещё одно заимствование из EAGLE-3 — перевзвешивание лоссов для разных драфтовых позиций. Ошибка в первых токенах обходится дороже, чем в поздних, поэтому ранним позициям задаётся больший вес. Коэффициент, управляющий перевзвешиванием, авторы выбирают примерно равным половине размера блока. Эксперименты показывают пользу такого подхода. Авторы сравнивают модели с различным количеством слоёв и рассматривают как среднюю длину принятой цепочки, так и итоговое ускорение. Пятислойная конфигурация оказывается наиболее удачной — увеличение глубины до восьми слоёв несколько улучшает acceptance length, однако сам драфтинг становится медленнее. Отдельно изучается количество слоёв основной модели, с которых берутся скрытые представления для формирования контекста. Здесь использование большего количества слоёв в целом оказывается полезным. Авторы также проверяют, как модель адаптируется к разным размерам блока в зависимости от того, с каким размером она обучалась. Результат ожидаемый: модель, обученная на большом блоке, относительно хорошо работает с меньшими блоками, тогда как модель, обученная на маленьком блоке, заметно хуже переносится на большие размеры. Ещё один эксперимент посвящен адаптации к длинному контексту. Авторы сравнивают базовую модель, обучавшуюся на контекстах до 4 тысяч токенов в течение шести эпох, с моделью, которую дополнительно обучали ещё три эпохи на контекстах длиной до 10 тысяч токенов. Результаты показывают, что дополнительное обучение на длинных контекстах улучшает работу модели с ними во время инференса. В последней абляции сравниваются обычная подача представлений на вход трансформерных слоёв и KV-injection. Эксперименты проводятся как для авторегрессионного, так и для диффузионного драфтинга. Способ внедрения через KV-проекции оказывается лучше обычной подачи на вход. DFlash показывает более высокое ускорение относительно EAGLE-3. Также авторы проводят сравнения с MTP. Однако здесь стоит обратить внимание, что сравнение с EAGLE нельзя назвать полностью честным из-за разных длин драфтовых цепочек и используемых кернелов: для EAGLE используется дерево глубины 7, а для собственного метода — цепочка длины 10. Применение более сложного и тяжеловесного кернела tree-attention для EAGLE не позволяет изолировать влияние метода на ускорение от эффективности верификации. Разбор подготовил ❣ Даниил Маслов Душный NLP