Ускорение диффузионных языковых моделей: метод MRP обещает до 1.42x прироста скорости
Группа исследователей представила новый метод Multi-token Residual Prediction (MRP), который позволяет существенно ускорить работу диффузионных языковых моделей (DLM). Эти модели генерируют текст, итеративно устраняя шум в маскированных последовательностях токенов, что даёт компромисс между параллелизмом и качеством по сравнению с авторегрессионными моделями.
В текущей практике число токенов, декодируемых за шаг, регулируется порогом уверенности, и качество монотонно ухудшается при увеличении числа токенов, обрабатываемых за один шаг. MRP решает эту проблему с помощью лёгкого модуля, который позволяет учитывать зависимости между токенами в рамках одного прямого прохода базовой модели.
Метод использует ключевое свойство процесса устранения шума: распределения логитов на соседних шагах оказываются очень похожими. Вместо того чтобы запускать модель второй раз для получения логитов следующего шага, MRP предсказывает остаток между шагами на основе скрытых состояний модели. Это позволяет эффективно обрабатывать больше токенов за один прямой проход.
MRP может работать в двух режимах: прямой декодинг (direct decoding), где скорректированные логиты используются без проверки для настройки компромисса между качеством и скоростью, и спекулятивный декодинг (speculative decoding), где предложения MRP проверяются базовой моделью для безошибочного ускорения.
Эксперименты проводились на моделях SDAR масштабом 1.7B, 4B и 8B параметров на задачах логического вывода и генерации кода. Результаты показали безошибочное ускорение до 1.42x в среде SGLang. Таким образом, MRP представляет собой эффективный способ повышения производительности диффузионных языковых моделей без потери качества генерации.
