Dynamic Chunking for End-to-End Hierarchical Sequence Modeling

меня упомянул Борис среди классных каналов, считаю своего рода успехом. более того, я был удостоен чести, что ко мне единственному придрались - по мне тоже ачивка ❤️ - за отсутствие заглавных букв
Ну, что ж..🤩

Вопрос токенизации для языковых (и не только) моделей стал более популярным после выхода Byte Latent Transformer (BLT), где токены заменялись на семантические “патчи”, которые разделялись предиктами отдельной модели (она предсказывала энтропию в авторегрессивном стиле) + SpaceByte использует “более натуральные” символы в английском для более семантичного разделения (e.g., /, :, ])

Но Альберт Гу, по его же словам, хотел чего-то более эстетичного и укладывающегося в end2end пайплайн - дифференцируемый чанкинг, который не только может быть применим и к языкам другой природы (китайский, код, днк), так еще и иметь интерпретируемость (и точно не содержать информацию о следующем байте)

И вот как сейчас выглядит версия H-Net с роутингом по аналогии МоЕ:
- Есть проекции q/k для каждого символа (чанка с предыдущей стадии), через которые определяется их схожесть
- Когда же между соседними репрезентациями наблюдается снижение в косинусном расстоянии (по трешхолду в 0.5), то считаем, что граница чанка определена
- Таким образом получаем и “токены” более осмысленно разграниченными, и длину последовательности сокращаем. При том за репрезентацию всего чанка берется первый эмбеддинг чанка (есть абляции на другие виды пулинга, результат не меняется)
- Далее основная сеть, которая суть заключает в себе миксеры на уровне последовательности по типу трансформер блока или мамбы2
- А дечанкинг (он же апсемплинг) дублирует репрезентацию чанка, пока не встретит новый, чтобы потом дублировать новый эмбеддинг с домножением на уверенность роутинг модуля

Все это приправляется EMA, STE и еще некоторыми добавлениями в архитектуру (доп нормализации и проекции, разные лернинг рейты под соответствующую глубину модулей) для стабилизации градиентного флоу в столь дискретном сетапе. Но, видимо, ничего не работает без load balancing как в МоЕ, который выставляет таргет количество символов, которое должно залетать в чанк (экспериментировали здесь с 6)

Первичные эксперименты выставили на длине в 8192 utf-8 символа по моделям в 760M & 1.3B + на китайском и коде проверили, сравнивали по FLOPS'ам с бейзлайнами (кстати с BLT нет результатов ибо они не захотели тратить время и компьют на трейн доп модели к основной). Визуализация границ чанков во многих местах выглядит интересно, часто попадает в нужные разделения, иногда граница лежит на какой-то хрени → не все так гладко и путь наименьшего сопротивления в обучении все еще ведет к костылям

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

Есть еще интересная секция про то, как дистиллить BPE бекбоны в чанкинг архитектуры, чтобы не полетели эмбеддинги, завязанные под токены - добавляют лоссы на матчинг векторов со своим начальным состоянием и пропускают начала токенов через сеть. Но пока непонятно

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

👀 paper, code

π∗0.6: a VLA That Learns From Experience

Давно меня не было😚😚😚

В последний раз про pi.website я писал, когда они сделали инпеинтинг для VLA с флоу матчингом (они уже успели и сделать его по классическому рецепту, обучаясь в inpainting-aware стиле). теперь же они над pi05 нафигачили то, что назвали pi06* (под звездочкой имеется в виду оффлайн рл)

Ну и скорее они надстроили не над pi05, а над pi06, где отличий немного (и мало деталей)

- гемма2 → гемма3 4B в качестве влм бекбона
- больше по параметрам флоу матчинг модуль для предикта действий
- больше датасет, качественнее промпты
- для эксперта действия теперь обрабатываются некаузальным аттеншном (и по идее тоже подаются в контекст вместе с языковыми и другими токенами)
- 5 шагов на инференсе вместо 10 в pi0 & pi05

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

- состоит эта value функция из меньшей влм с той же архитектурой, которая по обсервейшну и промпту выдает бины, чтобы предсказывать ретерны (сумму наград за часть эпизода, в данном случае это получается типа time-to-success)
- обучается она на многих тасках, чтобы потом через эту модель высчитывать advantage, который будет подаваться как условие в VLA
- VLA же обучается с этим advantage и всеми другими плюшками после SFT фазы + так же проводятся интервенции со стороны человека, чтобы чет поправить
- но не все так гладко: подаются не сами адвантаджи, а флаг, зависящий от трешхолда под каждую таску (превосходит ли адвантадж 30-ый или 40-ой перцентиль, в зависимости от того, какой этап обучения. кстати, в зависимости от этапа адвантадж тоже по-разному считается)

То есть, было: просто претрен на новых больших данных с бОльшей моделью (pi06)

Стало: претреним на мульти-таск сетапе value функцию → претрен на данных, где есть разметка advantage → дообучаем value под конкретную таску и затем дообучаем VLA на эту же таску с advantage разметкой → делаем роллаут моделей (и исправляем экспертными интервенциями, где надо) → дообучиваем value и VLA на обновленных данных (даггер стайл)

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


👀 paper, blog