Real-Time Execution of Action Chunking Flow Policies

снова pi

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

авторы же призадумались

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

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

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

помимо реальных примеров со складыванием вещей, посуды и проводами так же авторы много экспериментов сделали на кинетиксе, где сделали 12 бенчмарков (10 сред от оксфордских авторов и 2 свои, на которых данные собрали при помощи RPO для имитации робо траекторий)

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

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

👀 paper, code for kinetix, demo

A Stable Whitening Optimizer for Efficient Neural Network Training

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

а поконкретнее, главный автор углубился в недостатки шампуня, который аккумулирует multiplication products градиентов для их ортогонализации, что эмпирически помогает в обучении, и недавний muon можно интерпретировать как accumulation-free shampoo

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

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

главный автор аккуратно противопоставляет этому свой пропозал
- раз мы уже выяснили, что ортогонализация градиентов полезна, так может удастся в нужное место вставить еще и матричное разложение (второй вечный трюк в линале): будем хранить в кэше не инвертированные градиенты, а спектральное разложение от этих сущностей. интуитивно храним собственные вектора, которые соответствуют направлениям ковариации градиентов (hello pca) → более стабильные репрезентации
- и уже в eigen пространстве в качестве нормализации можно использовать просто взятие знака, что теоретически ограничивает спектральную норму и все получается гуд, так еще и гиперпараметр периода обновления кэша менее остро влияет на сходимость
- но (видимо без этого не заводится, cheap trick) надо добавлять EMA для параметров модели

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

все эксперименты проводились с джаксом на скейле в 160М трансформере по языковому моделированию + классификация картинок + диффузия, при том в сетапе упоминаются 32 карты TPUv3 (не особо понятно куда модель такого размера разворачивать на столько карт). есть так же и имплементация на пайторч, но на ней эксперименты не проводились (и в том числе нету понимания, останется ли торчовский код стабильным при раскатке на несколько карт)

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

👀 paper, code

Model-Preserving Adaptive Rounding

Альберт Тсенг может быть вам знаком по методам квантизаций qtip, quip/quip# и обучении ллм в mxfp4 , но не такому как quartet. он снова сделал квантизацию и получил алгоритм YAQA (Yet Another Quantization Algorithm)

GPTQ/LDLQ и AWQ методы производят квантизацию через прокси лосс разницы между активациями для отдельного слоя - layerwise mse + там присутствует гессиан для каждого слоя, который можно выразить через layer_input.T @ layer_input

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

используются те факты, что гессиан КЛ = fisher information matrix, которую можно эмпирически посчитать через gradient_loss.flatten() @ gradient_loss.flatten().T (один бекворд пасс) + произведение кронекера эквивалетно произведению с рангом 1, что можно получить через hessian-vector products, которые опять-таки хорошо компануются с бекворд пассом, упомянутым ранее, а следовательно и FSDP

вот так авторы и приближают оригинальный гессиан - через несколько итераций (до 3, power iterations) кронекер матрицами. при том они выводят 2 способа А и В
- А: дешевле (30 гпу-часов для 10В модели на 20М токенах), смещен, ниже разброс
- В: подороже (50 гпу-часов на 120М токенах), несмещен, но выше разброс → выше качество

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

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

по экспериментам - проверяются на лламе и гемме, где к yaqa используют квантизатор qtip и домножение на матрицы Адамара для сглаживания. по всем битрейтам примерно на треть деркзо бьет все, что есть. точнее - все, с чем сравнивались, ибо насчет PV-Tuning & AQLM есть вот такой не менее дерзкий комментарий

We do not directly compare to PV-Tuning since there are no public PV-Tuning models with either the QTIP or INT4 quantizer. However, LDLQ with the QTIP quantizer already outperforms PV-Tuning with the learnable AQLM quantizer on Llama 3.1, so we expect YAQA with QTIP to outperform PV-Tuning as well

👀 paper, code

Reinforcement Learning with Action Chunking

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

и вот сергей левин со своими студентами нацелился на применение этого трюка для классического пайплайна actor-critic q-learning’a. формулы обобщаются при том довольно интуитивно понятно и перестают быть марковскими - везде одно действие меняется на чанк действий, что даже улучшает понятие n-step return’a, где использовали для расчета значения критика n шагов вперед вместо одного, ибо тогда убирается смещение off-policy действий этого чанка действий

имплементится это через диффузию и флоу матчинг с ограничением на приверженность behavior policy в offline и offline-to-online рл сетапах. при том в сетапе с диффузией КЛ ограничение между политиками реализуют через best-of-n sampling (BFN), а с флоу матчингом сшивание идей происходит более гладко, без изменений в ключевых местах алгоритма FQL. экспы проводят над RLPD, где внутри онлайн степов батч состоит наполовину из онлайн и оффлайн буфферов

при том предикт по чанкам улучшает момент эксплорейшна, ибо, как спекулируют авторы, действия внутри одного чанка становятся более связанными относительно друг друга (в сравнении с 1-step методами) → при инференсе можем ожидать поведение получше + sample efficiency

пока и остается большой вопрос по поводу размера чанка, который сильно влияет на перформанс (на OGBench 10 действий в одном чанке у авторов лучше чем 25), а по балансу между рантаймом и sample efficiency неплохо было бы перепроверить, действительно ли обучение происходит быстрее бейзлайнов

👀 paper, code