Speculative Speculative Decoding – асинхронизируй это.
Вышел просто мёд для глаз моих в сфере оптимизации инференса.😜 Звать будем сокращённо SSD. Чуваки из TriaDao выпустили асинхронную версию speculative decoding.
Это важно прочитать до того как пойти дальше. Напомню, что стандартное спекулятивное декодирование ускоряет инференс LLM, используя быструю малую модель для предсказания токенов, которые затем проверяются большой моделью. Как это делается?
Пусть вероятность токенов из модели черновика q(x), вероятность этих же токенов для модели верификатора p(x), тогда:
1. Если q(x) <= p(x) мы принимаем драфт.
2. Иначе реджект с вероятностью p(x)/q(x).
3. Если отклонён, остаток добирается из распределнния разницы norm(max(0, p(x) — q(x)).
Логично, что для каких-то токенов там будет 0, далее остальные из хвоста топК будут давать положительную разницу. Её отнормируем и примем за вероятности и уже оттуда выберем.
Ну и типа это сродни x~p(x) с сходимостью по вероятности.
Однако SD само по себе имеет последовательное ограничение: черновая модель должна ждать завершения верификации перед началом следующей спекуляции.
SSD устраняет эту зависимость, запуская черновую модель на отдельном устройстве асинхронно – во время верификации она предсказывает вероятные исходы проверки и готовит спекуляции для них заранее. Если «хороший» исход оказывается в подготовленном наборе, спекуляция возвращается мгновенно, полностью устраняя задержку от малой модели.
Также вводится понятие saguaro sampling 💃— это новая схема сэмплирования, которая изменяет форму остаточного распределнния (такое вот а-ля температура) на наиболее вероятных токенах черновика, уменьшая их вероятность при сэмплировании.
Кажется, что тут есть противоречия, верно? Фраза прям витиеватая. Ща поясню на пальцах. Тут над вспомнить как в начале я писал про стратегию SD. Вернитесь перечитайте.
Итого, проблема, которую решаем. Нужно угадать, какой токен выберет большая модель после проверки.
Решение:
1. Берём топ-5 самых вероятных токенов по мнению маленькой модели.
2. Искусственно занижаем их вероятности при генерации черновика.
3. Это заставляет «остаточное распределение», куда большая модель смотрит, концентрироваться на этих же топ-5 токенах, без отбрасывания самых топовых, как я писал выше.
Почему это работает:
Большая модель смотрит: «А что если вычесть то, что уже предложил черновик?»
Если черновик «скромничает» и занижает вероятности топовых токенов → разница становится больше → большая модель с большей вероятностью выберет именно их.
Ещё случаются в таком семплинге промахи из-за кандидатов на асинхронных gpu, чем больше батч. Причём размер батча отсечки, когда какая стратегия, тоже приводят. Отсюда и fallback стратегии. Выкладки по стратегии fallback и оптимизации kv-cache, кому интересно, читайте статью. Формулы там повеселее, но приведу ещё один пример на пальцах.
Аналогия с рестораном.
Представь, что ты работаешь в ресторане и готовишь заказы наперёд. Твоя стратегия как в SSD:
Пока клиент ест — верификация идёт, ты готовишь топ5 возможных следующих блюд на выбор.
Кладёшь их на полку — kv кэш спекуляции.
Попадание в кэш. Клиент заказывает «пасту карбонара». У тебя есть на полке! → Мгновенно отдаёшь.
Промах кэша. Клиент заказывает «ризотто с грибами». На полке нет! → Приходится готовить с нуля, что медленно, но уже условно не N, а log(N) сложность. Тут предлагают выдать случайный токен или работать с вероятными n-gramms.
В плане рестика, конечно, не самый вариант дать челу рандомное блюдо, но тут и не ресторан 👍
В результате получаем до 2X быстрее, чем оптимизированное спекулятивное декодирование. До 5Х быстрее авторегрессионного декодирование. 😎
Всё. 📦
