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