LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

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

в чем основная суть? как оказывается, в основном, финальные токены для генерации формируются не на последнем слое, а где-то внутри модели с дальнейшей возможностью дальше поменяться, но такое происходит не во всех случаях (в среднем токену для полного формирования надо 23.45 слоев из 32-ух слойной лламы).

назвали это early exit - вместо того, чтобы ждать последнего слоя эмбеддингов, мы просто на вход ЛМ голове подаем эмбеддинги из промежуточных слоев (или просто слоев пораньше)

но есть небольшой питфолл в том, что голова научена только декодить эмбеддинги с последнего слоя → возникает необходимость дотренивать голову уметь декодить эмбеддинги с любого слоя

наивно это сделать не получается, чтобы работало + мы хотели бы инферить на как можно мЕньшем количестве слоев ⇒ чтобы токены генерились не на последних слоях. из этого прикрутили layer dropout с гипером, который экспоненциально увеличивается с глубиной сетки

да еще и забустили кроссэнтропию с поправкой на стимуляцию к ealry exit, где циркулярно или постепенно брали эмбеды из ранних слоев и осуществляли early exit

окей, раньше из сетки выходить мы поняли как, но как убедиться в том, что аккураси по сравнению со всей моделью не падает (падает незначительно)? - их навороченный спекулятивный декодинг (уже писали тута и тута), который они назвали self-speculative decoding

в этом декодинге они прогоняют всю модель для предикта следующего токена для каждого драфт токена, и все это в один форвард пасс + на самом деле надо прогонять не всю модель, а только оставшиеся после early exit’a (так еще и кв кэш сохраняется с этих драфт токенов). по экспам наблюдается спидап в 1.86 раз по заверениям авторов

выглядит красиво (и вроде практично) с дополнительными затратами в доп претрейне/файнтюнинге и подборе гиперов под дропаут и новый лосс вместе с лернинг рейтом и тд

код есть, кроме трейна, в принципе как и многого другого

👀LINK