Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models
seems like Альберт Гу стал частью нового стартапа cartesia.ai, которая продвигает on-device аудио и языковые модели (скорее всего, на основе ссмок)
ну и конечно, чтобы засунуть подобные модели в девайс, может понадобиться дистилляция. вот авторы и решили дистиллировать трансформеры в ссм - или более эпично - quadratic to sudquadratic mode
и таковое бы было менее удобно, если бы не недавняя ссд, которая позволяет свести и ссм, и аттеншн к матрикс миксерам (под общую формулу матричного умножения)
so, авторы перед тем как непосредственно дистиллировать phi-1 (хз почему выбор пал именно на нее), сравнили разные матрикс миксеры на аппроксимацию софтмакс аттеншна изолированно от обучения модели, то есть поставили отдельную задачу обучить разные матрикс миксеры аппроксимировать аттеншн мапы из разных слоев лламы 7б на небольших последовательностях из 4 датасетов → как оказывается, ссд фиттится лучше всех (для опять-таки ллмки относительно небольшого скейла)
дистилляцию еще решили разделить на 3 этапа
1. фит на матрикс миксеры - то есть сведение аутпутов между операциями аттеншна и ссд соответственно от инпутов трансформера. для этого пришлось немного переделать ссд из mutli-value в multi-head + отказаться от одной из активаций для лучшего сведения аттеншна и мамбы. хорошая новость заключается в том, что это можно распараллелить
2. фит между блоками - включаем теперь не только аттеншн и ссм, а так же леернормы, млп и остальное, что итеративно повторяется в архитектуре. етот этап так же параллелится
3. классическая end-to-end дистилляция. авторы так же упоминают, что здесь можно попробовать зафризить млп блоки и не потерять в качестве дистилляции. а это в свою очередь снизит количество обучаемых параметров наполовину
в итоге Phi-Mamba обучали по 80M/160M/2.76B токенов на каждом из этапов. по сравнению с Phi-1.5-1.3B которая обучалась на 150B токенах это очень круто. смущает разве что момент, что дистиллировали в 1.5B модель (pogchamp). для такого бы добавить аблации всякие, что такая дистилляция имеет смысл - хорошо, что данных намного меньше понадобилось, но легче ли это переносится чем трансформер или имеет меньше требований по памяти (скорее всего да, но опять-таки, неплохо бы посмотреть на наглядное сравнение)
имхо - появился какой-то интересный дым, теперь поближе бы подобраться к огню
👀LINK
seems like Альберт Гу стал частью нового стартапа cartesia.ai, которая продвигает on-device аудио и языковые модели (скорее всего, на основе ссмок)
ну и конечно, чтобы засунуть подобные модели в девайс, может понадобиться дистилляция. вот авторы и решили дистиллировать трансформеры в ссм - или более эпично - quadratic to sudquadratic mode
и таковое бы было менее удобно, если бы не недавняя ссд, которая позволяет свести и ссм, и аттеншн к матрикс миксерам (под общую формулу матричного умножения)
so, авторы перед тем как непосредственно дистиллировать phi-1 (хз почему выбор пал именно на нее), сравнили разные матрикс миксеры на аппроксимацию софтмакс аттеншна изолированно от обучения модели, то есть поставили отдельную задачу обучить разные матрикс миксеры аппроксимировать аттеншн мапы из разных слоев лламы 7б на небольших последовательностях из 4 датасетов → как оказывается, ссд фиттится лучше всех (для опять-таки ллмки относительно небольшого скейла)
дистилляцию еще решили разделить на 3 этапа
1. фит на матрикс миксеры - то есть сведение аутпутов между операциями аттеншна и ссд соответственно от инпутов трансформера. для этого пришлось немного переделать ссд из mutli-value в multi-head + отказаться от одной из активаций для лучшего сведения аттеншна и мамбы. хорошая новость заключается в том, что это можно распараллелить
2. фит между блоками - включаем теперь не только аттеншн и ссм, а так же леернормы, млп и остальное, что итеративно повторяется в архитектуре. етот этап так же параллелится
3. классическая end-to-end дистилляция. авторы так же упоминают, что здесь можно попробовать зафризить млп блоки и не потерять в качестве дистилляции. а это в свою очередь снизит количество обучаемых параметров наполовину
в итоге Phi-Mamba обучали по 80M/160M/2.76B токенов на каждом из этапов. по сравнению с Phi-1.5-1.3B которая обучалась на 150B токенах это очень круто. смущает разве что момент, что дистиллировали в 1.5B модель (pogchamp). для такого бы добавить аблации всякие, что такая дистилляция имеет смысл - хорошо, что данных намного меньше понадобилось, но легче ли это переносится чем трансформер или имеет меньше требований по памяти (скорее всего да, но опять-таки, неплохо бы посмотреть на наглядное сравнение)
имхо - появился какой-то интересный дым, теперь поближе бы подобраться к огню
👀LINK
Retrieval Head Mechanistically Explains Long-Context Factuality
мы уже говорили о том, что внутри трансформеров можно находить интерпретируемы однозначные аспекты - будь то направление в пространстве параметров или индуктивные головы, которые так важны для ин-контекст лернинга
и в этот раз авторы смогли найти в аттеншне головы, которые отвечают за ретривал функцию. составляют они малую часть от общего количества голов (что логично), при том, если их отключить, то пропадут не только ретривал способности у модели, но так же и CoT способность (ибо для того, чтобы сохранить цеопочку рассуждений, надо уметь опираться на какие-то факты из прошлого)
к тому же если еще и продолжать тренить на удлинение контекста, то ретривал головы не меняют своего назначения (и точно так же не добавляется новых голов к этому подмножеству)
смущает разве что их мнение по поводу компрессии кв кэша - поскольку ретривал головы не могут без аттенда к предыдущим токенам, то им позарез нужен кв кэш, в то время как (якобы) остальные головы могут пренебречь кв кэшем и использовать только intrinsic knowledge из весов. имхо это частино верно (либо я неправильно понял мысль авторов) - тот факт, что мы не можем найти однозначно интерпретируемые фичи, за которые отвечают другие головы аттеншна и которые связаны с историей контекста, не означает, что этим головам в принципе не нужен контекст
а значит вопрос, который ставят авторы, может звучать примерно как-то так - возможно ли сделать оптимальный роутинг (по аналогии как здесь) относительно голов каждого слоя, следует ли выделять кэш
👀LINK
мы уже говорили о том, что внутри трансформеров можно находить интерпретируемы однозначные аспекты - будь то направление в пространстве параметров или индуктивные головы, которые так важны для ин-контекст лернинга
и в этот раз авторы смогли найти в аттеншне головы, которые отвечают за ретривал функцию. составляют они малую часть от общего количества голов (что логично), при том, если их отключить, то пропадут не только ретривал способности у модели, но так же и CoT способность (ибо для того, чтобы сохранить цеопочку рассуждений, надо уметь опираться на какие-то факты из прошлого)
к тому же если еще и продолжать тренить на удлинение контекста, то ретривал головы не меняют своего назначения (и точно так же не добавляется новых голов к этому подмножеству)
смущает разве что их мнение по поводу компрессии кв кэша - поскольку ретривал головы не могут без аттенда к предыдущим токенам, то им позарез нужен кв кэш, в то время как (якобы) остальные головы могут пренебречь кв кэшем и использовать только intrinsic knowledge из весов. имхо это частино верно (либо я неправильно понял мысль авторов) - тот факт, что мы не можем найти однозначно интерпретируемые фичи, за которые отвечают другие головы аттеншна и которые связаны с историей контекста, не означает, что этим головам в принципе не нужен контекст
а значит вопрос, который ставят авторы, может звучать примерно как-то так - возможно ли сделать оптимальный роутинг (по аналогии как здесь) относительно голов каждого слоя, следует ли выделять кэш
👀LINK
👍3
DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads
автор недавно упомянутых ретривал голов и аттеншн синков решил объединить эти 2 концепции и нехило так сэкономить по памяти для кв кэша
продолжая идею, которой они прогревали в статье про ретривал головы, про компрессию кв кэша - давайте хранить весь кэш только для этих голов, а остальным оставлять кэш только для синк токенов и недавних токенов (степень “недавности” определяется гипером в соответствии с трейдоффами по памяти и перформансу)
определяют же ретривал головы в этот раз по другому сетапу - ставят обучаемый гейт на каждую голову каждого слоя, с помощью которого фиксируют пропорцию между фулл аттеншном и “стриминг аттеншном” + л1 регуляризация. получается num_layers x num_heads параметров, которые можно зафиттить на классический форвард пасс за 2к степов на 8 а100 (по заверениям авторов). обучались же на их синтетике (не особо понятно почему это было проще чем запустить сразу на needle in a haystack + есть вопросы по поводу того, действительно ретривал головы не меняются если менять датасеты на которых мы их ищем)
а в принципе звучит даже интересно, посмотрим как приживется, если они код смогут попроще причесать
👀LINK
автор недавно упомянутых ретривал голов и аттеншн синков решил объединить эти 2 концепции и нехило так сэкономить по памяти для кв кэша
продолжая идею, которой они прогревали в статье про ретривал головы, про компрессию кв кэша - давайте хранить весь кэш только для этих голов, а остальным оставлять кэш только для синк токенов и недавних токенов (степень “недавности” определяется гипером в соответствии с трейдоффами по памяти и перформансу)
определяют же ретривал головы в этот раз по другому сетапу - ставят обучаемый гейт на каждую голову каждого слоя, с помощью которого фиксируют пропорцию между фулл аттеншном и “стриминг аттеншном” + л1 регуляризация. получается num_layers x num_heads параметров, которые можно зафиттить на классический форвард пасс за 2к степов на 8 а100 (по заверениям авторов). обучались же на их синтетике (не особо понятно почему это было проще чем запустить сразу на needle in a haystack + есть вопросы по поводу того, действительно ретривал головы не меняются если менять датасеты на которых мы их ищем)
а в принципе звучит даже интересно, посмотрим как приживется, если они код смогут попроще причесать
👀LINK
👍6
Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs
Сегодня смотрим реальный трансформер для молекул c SE(3)/E(3) эквивариатностью☯️ Equiformer - разработка группы Atomic Architects из MIT ⚡️ Очень крутые ребята! Работает он на equivariant модулях с irreducible representations от сферических функций (обозначаются как type-L вектора, где L - степень репрезентаций). Такое в целом мы уже видели в других работах, но тут красота статьи в деталях
Сначала рассмотрим любимый аттеншен
Фичи с source и target ноды подают в линейный слой, а потом складывают и умножают тензорно на эмбеддинги сферической функции относительной позиции атомов. Это операция сохраняет семантические и геометрические признаки. Для аттеншена (весов) юзают только скалярный вектора так как веса аттеншена инвариантны. Вместо дот продакта берут MLP как в GATv2 для более гибких паттернов. Сами фичи после преобразований нод разбивают на скаляры и на другие. По скалярам считают веса, а по другим делают фичи вершин через эквивариантную активацию, тензорное умножение и линейный слой. Потом перезвешивают и все это можно делать как MHA
Еще есть другие мелкие детали касательно разных слоев
1) Линейные слои обрабатывают различные type-L вектора отдельно + не используют сдвиг, так как они сломают эквивариантность при L > 0
2) LN нормализуется как RMS L2-нормы type-L вектора. Также нет сдвигов и средних для L > 0
3) Используют более эффективное тензороное умножение, где каждый тип output вектора зависит от одного типа input вектора. Но вообще не понятно, почему на визуализациях type-0 зависит только от type-0 и type-1, но не от type-2
На экспах классно показывает себя на OC20, а вот на QM9 средне, но в целом это маленький датасет, поэтому пофиг
👀 LINK
Сегодня смотрим реальный трансформер для молекул c SE(3)/E(3) эквивариатностью
Сначала рассмотрим любимый аттеншен
Фичи с source и target ноды подают в линейный слой, а потом складывают и умножают тензорно на эмбеддинги сферической функции относительной позиции атомов. Это операция сохраняет семантические и геометрические признаки. Для аттеншена (весов) юзают только скалярный вектора так как веса аттеншена инвариантны. Вместо дот продакта берут MLP как в GATv2 для более гибких паттернов. Сами фичи после преобразований нод разбивают на скаляры и на другие. По скалярам считают веса, а по другим делают фичи вершин через эквивариантную активацию, тензорное умножение и линейный слой. Потом перезвешивают и все это можно делать как MHA
Еще есть другие мелкие детали касательно разных слоев
1) Линейные слои обрабатывают различные type-L вектора отдельно + не используют сдвиг, так как они сломают эквивариантность при L > 0
2) LN нормализуется как RMS L2-нормы type-L вектора. Также нет сдвигов и средних для L > 0
3) Используют более эффективное тензороное умножение, где каждый тип output вектора зависит от одного типа input вектора. Но вообще не понятно, почему на визуализациях type-0 зависит только от type-0 и type-1, но не от type-2
На экспах классно показывает себя на OC20, а вот на QM9 средне, но в целом это маленький датасет, поэтому пофиг
Please open Telegram to view this post
VIEW IN TELEGRAM
Please open Telegram to view this post
VIEW IN TELEGRAM
❤6
Cut Your Losses in Large-Vocabulary Language Models
авторы вспоминают вычисление кросс энтропии
где e - финальные эмбеддинги, с - матрица, которая переводит из эмбеддингов в vocab_size вектор. и эта самая матрица
и чтобы снизить память, авторы сделали свой тритоновский кернел, который не засовывает сразу все логиты в глобальную память, а вычисляет логиты для верных токенов (во flash memory из-за концепции teacher forcing) и log-sum-exp на лету (по аналогии с онлайн софтмаксом из флеш аттн), т.е. не аллоцирует память на
к тому же еще при бекворде они заметили, что меньше чем 0.02% от нормализованных логитов являются ненулевыми (по большей части из-за действительно большого вокаб сайза) → не вычисляют градиент для элементов, которые не проходят трешхолд в 2е-12 для бф16 точности, что тоже снижает по памяти и прибавляет по скорости
оверолл голова ллмок начала теперь кушать по памяти в среднем гигабайт вместо 28 при батч сайзе в 65к (при том экспы ставили на модельках от 1.3B do 70B за что респект). есть правда вопросы насколько стабильно будет этот метод работать для претрена с нуля, ибо авторы только “файнтюнили”
очень понятно описано решение их кернела, как и сам код, рекомендуем к прочтению
выглядит интересно и прикольно, на последней картинке только с осторожностью относился бы к абсолютным числам, которые у них получились на экспах по замеру времени
👀LINK
авторы вспоминают вычисление кросс энтропии
loss = F.cross_entropy((e @ c.T).float(), targets)
где e - финальные эмбеддинги, с - матрица, которая переводит из эмбеддингов в vocab_size вектор. и эта самая матрица
e @ c.T, как оказывается, кушает оочень многоо памяти, а при увеличении вокаб сайза затрат по памяти становится еще большеи чтобы снизить память, авторы сделали свой тритоновский кернел, который не засовывает сразу все логиты в глобальную память, а вычисляет логиты для верных токенов (во flash memory из-за концепции teacher forcing) и log-sum-exp на лету (по аналогии с онлайн софтмаксом из флеш аттн), т.е. не аллоцирует память на
e @ c.T
к тому же еще при бекворде они заметили, что меньше чем 0.02% от нормализованных логитов являются ненулевыми (по большей части из-за действительно большого вокаб сайза) → не вычисляют градиент для элементов, которые не проходят трешхолд в 2е-12 для бф16 точности, что тоже снижает по памяти и прибавляет по скорости
оверолл голова ллмок начала теперь кушать по памяти в среднем гигабайт вместо 28 при батч сайзе в 65к (при том экспы ставили на модельках от 1.3B do 70B за что респект). есть правда вопросы насколько стабильно будет этот метод работать для претрена с нуля, ибо авторы только “файнтюнили”
очень понятно описано решение их кернела, как и сам код, рекомендуем к прочтению
выглядит интересно и прикольно, на последней картинке только с осторожностью относился бы к абсолютным числам, которые у них получились на экспах по замеру времени
👀LINK
👍6❤1
nGPT: Normalized Transformer with Representation Learning on the Hypersphere
в связи с последним issue, было бы интересно сделать пост
относительно известный оптимизатор Ilya Loshchilov решил соптимизировать архитектуру гпт. а именно он убрал классические нормализации и добавил нормализасион по одной и той же гиперсфере (о Боже неевклидова геометрия пошла) + добавили обучаемый скейлинг фактор (ибо без него не заработало вообще никак судя по всему)
так еще к тому же авторы вспоминают, что трансформер на самом деле внутри мета-оптимайзит подаваемые на вход функции или что-то типа того (может делать итеративно градиентный спуск и TD в аттеншне). потому они решают напрямую добавить обучаемые лернинг рейты в трансформер блок (названные эйгенлернинг рейтами), которые помогали бы более явно воспроизводить этот процесс мета-оптимизации
ну как будто что-то там действительно работает лучше - хоть один форвард по себе занимает намного больше времени (что так-то большое упущение для практических целей, но lucidrains смог подускорить), общее количество степов нужно меньше, при чем настолько, что он и по времени сходится быстрее обычной гпт (на малюсеньком скейле 0.5-1B параметров с не бОльшим датасетом)
ну а если чуть-чуть глубже копнуть по данному issue и ответе авторов, то очень-очень начинают терзать смутные сомнения, а как верить резам, когда якобы internal code не совпадает с опенсурсом внутри форварда трансформер блока (то есть мы должны поверить авторам, что для опенсурса они решили все с нуля переписать и совершить такую ошибку, нежели копипастнуть). но про презумцию надо помнить и потому верим😊 😊 😊 😊
👀LINK
в связи с последним issue, было бы интересно сделать пост
относительно известный оптимизатор Ilya Loshchilov решил соптимизировать архитектуру гпт. а именно он убрал классические нормализации и добавил нормализасион по одной и той же гиперсфере (о Боже неевклидова геометрия пошла) + добавили обучаемый скейлинг фактор (ибо без него не заработало вообще никак судя по всему)
так еще к тому же авторы вспоминают, что трансформер на самом деле внутри мета-оптимайзит подаваемые на вход функции или что-то типа того (может делать итеративно градиентный спуск и TD в аттеншне). потому они решают напрямую добавить обучаемые лернинг рейты в трансформер блок (названные эйгенлернинг рейтами), которые помогали бы более явно воспроизводить этот процесс мета-оптимизации
ну как будто что-то там действительно работает лучше - хоть один форвард по себе занимает намного больше времени (что так-то большое упущение для практических целей, но lucidrains смог подускорить), общее количество степов нужно меньше, при чем настолько, что он и по времени сходится быстрее обычной гпт (на малюсеньком скейле 0.5-1B параметров с не бОльшим датасетом)
ну а если чуть-чуть глубже копнуть по данному issue и ответе авторов, то очень-очень начинают терзать смутные сомнения, а как верить резам, когда якобы internal code не совпадает с опенсурсом внутри форварда трансформер блока (то есть мы должны поверить авторам, что для опенсурса они решили все с нуля переписать и совершить такую ошибку, нежели копипастнуть). но про презумцию надо помнить и потому верим
👀LINK
Please open Telegram to view this post
VIEW IN TELEGRAM
Please open Telegram to view this post
VIEW IN TELEGRAM
👍4❤2
ADOPT: Modified Adam Can Converge with Any β2 with the Optimal Rate
на определенном этапе заведения модели, которая не заводится, начинаешь задумываться про гиперы, которые стоят за оптимизатором (помимо лернинг рейта) и самого оптимизатора. например, на беты
и как мы уже упоминали, в то время как первая бета отвечает за сохранение градиентов для первого момента, вторая бета отвечает за сохранение истории в бегущем среднем вторых моментов градиента (что логично). и с точки зрения теории адам (да и в принципе все адаптивные методы) довольно плохо сходится, если только не выбирать эту вторую бету в зависимости от поставленной таски
но вот авторы-японцы (возможно) смогли это исправить и нескромно назвали метод ADaptive gradient method with the OPTimal convergence rate
и вот для того, чтобы вторая бета не имела такой сильный импакт на сходимость, они к удивлению меняют расчет первого момента - дополнительно делят градиент на данном таймстепе на корень из второго момента. простенько, со вкусом, достаточно нетривиально для данной специфики
по экспам где-то даже резы лучше достигаются - в том числе и на 7б лламе прогоняли (правда только ммлу, как любит замечать наш дорогой друг, без алаймент бенчмарков это не особо релевантно) + для мниста и цифара брали только резнет-18 но допууууустим
к тому же тут есть тоже предположение в их теории - о том что второй момент градиентов ограничен (менее сильное предположение в сравнении с предыдущим о том, что первый момент тож ограничен)
позабавило еще то, что в вывод в конце они зачем-то решили вставить проблему социального импакта мл алгоритмов (хотя статья чисто про оптимизатор)
а код оч классный, челики в сурс коде торча знатно так разбираются
👀LINK
на определенном этапе заведения модели, которая не заводится, начинаешь задумываться про гиперы, которые стоят за оптимизатором (помимо лернинг рейта) и самого оптимизатора. например, на беты
и как мы уже упоминали, в то время как первая бета отвечает за сохранение градиентов для первого момента, вторая бета отвечает за сохранение истории в бегущем среднем вторых моментов градиента (что логично). и с точки зрения теории адам (да и в принципе все адаптивные методы) довольно плохо сходится, если только не выбирать эту вторую бету в зависимости от поставленной таски
но вот авторы-японцы (возможно) смогли это исправить и нескромно назвали метод ADaptive gradient method with the OPTimal convergence rate
и вот для того, чтобы вторая бета не имела такой сильный импакт на сходимость, они к удивлению меняют расчет первого момента - дополнительно делят градиент на данном таймстепе на корень из второго момента. простенько, со вкусом, достаточно нетривиально для данной специфики
по экспам где-то даже резы лучше достигаются - в том числе и на 7б лламе прогоняли (правда только ммлу, как любит замечать наш дорогой друг, без алаймент бенчмарков это не особо релевантно) + для мниста и цифара брали только резнет-18 но допууууустим
к тому же тут есть тоже предположение в их теории - о том что второй момент градиентов ограничен (менее сильное предположение в сравнении с предыдущим о том, что первый момент тож ограничен)
позабавило еще то, что в вывод в конце они зачем-то решили вставить проблему социального импакта мл алгоритмов (хотя статья чисто про оптимизатор)
а код оч классный, челики в сурс коде торча знатно так разбираются
👀LINK