На неделе у нас был очень классный reading club по статье Coupling without Communication and Drafter-Invariant Speculative Decoding, из которой хочется поделиться одной интересной концепцией.
Но для начала стоит сказать про Gumbel-Max trick – способ семплировать из категориального распределения. Стандартный путь такой – считаем softmax над всем словарём и семплируем сам токен. Обе процедуры могут быть достаточно дорогими. Gumbel-Max делает следующее: к логитам прибавляем шум из распределения Gumbel(0, 1) и просто берём argmax. Математически это эквивалентно семплированию из softmax распределения, но при этом:
- softmax можно не считать вообще
- шум не зависит от логитов, его можно подготовить заранее, пока модель делает forward pass
- из коробки получается top-k семплирование (Stochastic Beams and Where to Find Them)
Теперь про идею Gumbel Coupling и ее применение для спекулятивного декодинга. Идея в том, чтобы использовать один и тот же gumbel шум при семплировании из драфтового и таргетного распределений. В таком случае можно записать корректный алгоритм с оценкой на acceptance rate, который будет просто сравнивать два токена между собой: токен от draft и target моделей. И вот тут получаются две интересные истории:
1. Не нужно материализовывать драфтовые логиты. В классической схеме верификации rejection sampling работает с обоими распределениями p и q одновременно, то есть драфтовые вероятности надо тащить через всю фазу верификации. С Gumbel coupling решение можно принимать, имея на руках только argmax от драфтера.
2. Воспроизводимость генерации. Вот это прям супер прикольно! В стандартной схеме итоговый сэмпл зависит от обоих распределений: меняется драфтер – меняется и то, что итоговая модель насемплит, даже при фиксированном seed. Получается, что чисто техническая оптимизация инференса влияет на то, что выходит из модели. С Gumbel coupling это свойство восстанавливается: при фиксированном seed выход не зависит от того, какой драфтер используется (и используется ли он вообще). Естественно, при условии, что все остальные источники недетерминизма мы победили (Defeating Nondeterminism in LLM Inference).
В презентации есть ещё иллюстрации к самому coupling-у, разбор алгоритма в vLLM и отсылки в сторону Speculative Speculative Decoding.
Но для начала стоит сказать про Gumbel-Max trick – способ семплировать из категориального распределения. Стандартный путь такой – считаем softmax над всем словарём и семплируем сам токен. Обе процедуры могут быть достаточно дорогими. Gumbel-Max делает следующее: к логитам прибавляем шум из распределения Gumbel(0, 1) и просто берём argmax. Математически это эквивалентно семплированию из softmax распределения, но при этом:
- softmax можно не считать вообще
- шум не зависит от логитов, его можно подготовить заранее, пока модель делает forward pass
- из коробки получается top-k семплирование (Stochastic Beams and Where to Find Them)
Теперь про идею Gumbel Coupling и ее применение для спекулятивного декодинга. Идея в том, чтобы использовать один и тот же gumbel шум при семплировании из драфтового и таргетного распределений. В таком случае можно записать корректный алгоритм с оценкой на acceptance rate, который будет просто сравнивать два токена между собой: токен от draft и target моделей. И вот тут получаются две интересные истории:
1. Не нужно материализовывать драфтовые логиты. В классической схеме верификации rejection sampling работает с обоими распределениями p и q одновременно, то есть драфтовые вероятности надо тащить через всю фазу верификации. С Gumbel coupling решение можно принимать, имея на руках только argmax от драфтера.
2. Воспроизводимость генерации. Вот это прям супер прикольно! В стандартной схеме итоговый сэмпл зависит от обоих распределений: меняется драфтер – меняется и то, что итоговая модель насемплит, даже при фиксированном seed. Получается, что чисто техническая оптимизация инференса влияет на то, что выходит из модели. С Gumbel coupling это свойство восстанавливается: при фиксированном seed выход не зависит от того, какой драфтер используется (и используется ли он вообще). Естественно, при условии, что все остальные источники недетерминизма мы победили (Defeating Nondeterminism in LLM Inference).
В презентации есть ещё иллюстрации к самому coupling-у, разбор алгоритма в vLLM и отсылки в сторону Speculative Speculative Decoding.