Forgetting Transformer: Softmax Attention with a Forget Gate

авторы из Милы и дипмаинда сделали очень естественную концепцию форгет гейта для софтмакс аттеншна

а естественность эта заключается в том, что

- если вспомнить линейный аттеншн, то его преимущество на инференсе (теоретически) заключается в возможности записать его в рекуррентном виде
- где уже дальше можно внести момент data-dependent форгет гейта (взятый по аналогии с LSTM или xLSTM если поновее (кстати его смогли отскейлить до 7б), только там экспоненциальное забывание)
- однако, это самое забывание можно внести и в привычную формулу параллельного линейного аттеншна (что в принципе похоже на GLA & обобщение ссм и аттеншна мамбы 2)
- но тогда там в формуле появляется уравнение forget_matrix * sim(q, k), которую снова можно развернуть в forget_matrix * softmax(q @ k)

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

softmax(q @ k + log[forget_matrix])

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

так еще и реализовано во флеш-аттн лайк манере, где эта матрица не высчитывается сразу, а постепенно, храня в HBM промежуточные результаты forget gate’ов для каждого таймстепа,и попарно высчитывается далее уже в онлайн стиле (в статье очень хорошо и понятно описано, рекомендуем)

сделали все на основе флаг аттеншн - оч крутая репа, да и сам код от FoX (Forgetting Transformer) не отстает

по состоятельности экспериментов можно сказать то, что реально присутствует улучшение в языковом моделировании, зеро-шот, в длинный контекст и все такое, однако скейл моделей был примерно в районе 125М-350М-760М, ну и скейл датасета соответствующий

👀LINK

Isaac GR00T N1: An Open Foundation Model for Humanoid Robots

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

у них здесь аналогично есть иерархичное разделение на систему 1 и систему 2 для того, что они называют ризонингом (ибо на хайпе), и совершением действия на 120 Гц. кстати как и у pi0 by pi.website, в качестве головы для действий выступает флоу матчинг голова (где таймстепы точно так же семплируются из бета распределения, интересно проводили ли какие-то махинации ли они с токенизацией)

моделька называется грут 2.2B (а сколько пи0 составляет неизвестно UPD: известно, спасибо комментаторам, примерно столько же), у которого VLM бэкбон это Eagle2) опять-таки от нвидиа. при этом есть и такие занимательные моменты

- на инференсе они делают 4 шага флоу матчинга (вместо 10 как у беркли/физикал интеллиженс)
- система 2 (VLM) принимает на вход картинки разрешения 224х224, что уже многолетняя классика, качества которой достаточно для робо задач
- эмбеддинги VLM они берут не последние, а из середины модели (12-ый слой, про такую технику мы уже говорили), может намекать о том, что языковые инструкции для роботов не так уж сложны для обработки (либо еще что-то непонятное что требует больше ресерча)
- еще так же во время обучения авторы добавили доп лосс на детекцию objects of interest, который реализован просто через отдельную голову для 2д координат (сделано для model spatial understanding)

по компьюту же авторы использовали вплоть до 1024 карточек (необязательно для обучения, но и для сбора данных). грут 2б же скушал примерно 50к гпу часов на h100 на претрейне, потом еще файнтюнили на А6000

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

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

но и этого мало! от этих видосов также нагенерили много синтетики, которые повышают диверсити траекторий действий. при том нагенерили в сценарии имг2видео, для чего использовали WAN 14B, а уже для составления робо траекторий выступала модель DexMimicGen (тоже от нвидиа)

We generate a total of around 827 hours of videos; it takes 2 minutes to generate a one-second video on an L40 GPU, and required approximately 105k L40 GPU hours (∼1.5 days) on 3,600 L40 GPUs.

короче говоря, вышло очень дорого на всех этапах, а демки залипательные

👀 link, code

Cosmos World Foundation Model Platform for Physical AI

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

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

но на чем учить? хз, видимо другой опции как 10^8 + 10^7 клипов (для претрена и файнтюна соответственно) у авторов не было. а получили они их из 20 млн видосов разного качества (от 720р до 4к), сами же клипы получились долгим путем обработки

- scene segmentation
- фильтрация на движения, качество, отсутствие перекрывающего текста и тд
- аннотация при помощи VLM
- deduplication

по категориям - от вождения, человеческого движения и активности, nature dynamics до вида от первого лица, динамического движения камеры и тп. собрали все примерно сбалансированно

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

- с дискретными токенами (на основе FSQ)
- и непрерывными (на основе ванильного автоэнкодера)
- при том везде в архитектурах есть вейвлет преобразование с темпоральными каузальными свертками и аттеншном

при том виды таких квантизаций, видимо, выбраны таким образом, чтобы не добавлять дополнительных слагаемых в финальный лосс при обучении моделей (ни кл терма например в случае VAE или коммитмент лосса в случае дефолтного VQ-VAE для непрерывного и дискретного случаев соответственно). а обучались эти токенизаторы на reconstruction loss + еще авторы добавили лоссы на optical flow & матрицу грама

для диффузии (помимо того что есть на картинках)

- непрерывные токены (хотя дискретные присутствуют и при обучении)
- не флоу матчинг, а диффьюжн скор (как из EDM) для DiT
- prompt upsampler
- mixed precision
- fsdp + context parallelism

декодер-онли трансформеры

- обучаемый softmax scaling factor
- z-loss как дополнительный терм чтобы логиты по магнитудам были не слишком большими
- tensor parallelism + sequence parallelism
- medusa на инференсе как вид спекулятивного декодинга

We train all of the WFM models reported in the paper using a cluster of 10,000 NVIDIA H100 GPUs in a time span of three months.

так же авторы еще пробовали файнтюнить

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

помимо метрик типа ошибки сэмпсона, LPIPS или 3д консистенси авторы еще решили прочекать как хорошо модели понимают физику (потому что вся вот эта космическая суета в принципе под эгидой Physical AI развивается) в сравнении и isaac sim в плане падения и движения предметов

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

так же второй из трех подпроектов космоса основывается на этих моделях - Cosmos-Transfer1, где происходит sim2real генерация из разных модальностей (сегментация, карта глубины и edge maps) в видео реальности

👀 link, code