Оптимизатор FARO снизил валидационные потери GPT на 0,0136 нат на 6 из 6 сидов

Современные рецепты обучения больших языковых моделей делят работу между оптимизаторами вручную: ортогонализованный импульс вроде Muon обновляет матрицы внимания и MLP, правила в стиле Adam отвечают за остальное, а на всю модель действует один общий бюджет обновления. Исследователи из Гонконгского университета науки и технологий (Гуанчжоу) предложили FARO — оптимизатор, который принимает эти решения по измерениям на самой обучаемой модели. На GPT-подобной модели, обученной на FineWeb, его вариант FARO-Adaptive (только выбор геометрии обновления) достиг потерь 3,2780 нат, а надстройка FARO-ERC (перераспределение бюджета по принципу равного вклада в риск, equal risk contribution) снизила потери на всех 6 сидах, то есть запусках с разной случайной инициализацией, в среднем на 0,0136 нат. Нат — естественная единица измерения кросс-энтропийных потерь. Код опубликован на GitHub.

Геометрия по результату, а не по типу слоя

Стандартный рецепт опирается на два допущения: тип слоя определяет подходящую геометрию обновления (матрицам — спектральная норма, векторам — покоординатная адаптивность), а одного скалярного бюджета достаточно для всей сети. Оба допущения взяты из человеческих эвристик и небольших сеток абляций, а не из измерений на обучаемой модели. Новые матричные оптимизаторы — NorMuon, Scion, Dion2, ARO, PolarGrad, Newton–Muon — меняют лишь одну деталь рецепта и по-прежнему фиксируют одну геометрию для всех матриц, которых касаются.

FARO берёт сам выбор за объект обучения. Он разбивает параметры на 117 блоков. Для каждого блока все кандидатные правила применяются к отложенной контрольной выборке, и остаётся то, что сильнее снижает потери. Пока ранжирование информативно, маршруты могут меняться, затем они замораживаются. Это ограничивает стоимость поиска и убирает выбор маршрутов из источников разброса. FARO — каркас, а не новое правило обновления: Muon, NorMuon и Adam-подобные правила входят в него как кандидаты, а не как конкуренты. В приложении авторы показывают, что поблочное решение согласуется с совместным по всем блокам при мягком допущении о раздельности, и оценивают границу связи, которую оно игнорирует.

В итоге FARO-Adaptive достиг 3,2780 нат на полном расписании обучения. Лучший ручной рецепт NorMuonAndAdam показал 3,2794. Система воспроизвела человеческую эвристику без ручной разметки геометрии по слоям.

Risk parity из портфельной теории

Вторая часть работы касается бюджета. После заморозки маршрутов у каждого блока остаётся свой множитель, масштабирующий обновление. Три скалярных контроллера — дробный множитель Келли, его непрерывная релаксация и правило с оценкой сложности — масштабируют один общий бюджет. Все три не отличились от риск-нейтрального FARO-Adaptive, что согласуется с наблюдением: расширение шага в погоне за риском непродуктивно.

Авторы объясняют это структурно. Скаляр умножает все бюджеты на общий коэффициент и движется только по диагонали симплекса бюджетов, поэтому не может решить, каким блокам нужно больше обновлений, а каким меньше. Поэтому они перераспределяют бюджет между блоками при неизменной сумме по правилу из портфельной теории — равному вкладу в риск (ERC, он же risk parity). В портфеле веса подбирают так, чтобы каждый актив давал одинаковую долю совокупного риска. В FARO блок играет роль актива, а отношение обновления к норме весов — роль волатильности.

Сырой риск по блокам различается на порядки. Поэтому его ограничивают (винзоризация) относительно медианы по блокам, множитель ограничивают сверху, а масштаб, сохраняющий суммарный бюджет с численной точностью, находят бисекцией. Перераспределение использует риск только предыдущего шага, так что текущий градиент не влияет на собственный бюджет. Бисекция даёт точное решение раздельной задачи с риском, где связь по бюджету несёт одна скалярная цена, — равные вклады в риск получаются из решения, а не навязываются. Благодаря этому сравнение с FARO-Adaptive отделяет перераспределение от простого изменения эффективного learning rate.

Скромный, но устойчивый выигрыш

FARO-ERC снизил потери на чистых данных на всех 6 сидах, в среднем на 0,0136 нат при том же объёме вычислений. Тест знаков, парный тест и доверительный интервал согласуются между собой, а эффект превышает измеренный разброс между запусками. Время и память не растут. Механизм такой: ERC урезает самые рискованные блоки, а не расширяет самые безопасные. Авторы сами называют выигрыш скромным, но устойчивым.

Метод проверили и на искажённом потоке данных: средние потери всё равно снизились, а большинство сидов осталось за FARO-ERC.

Практический вывод для тех, кто настраивает претрейн: ручное разделение «матрицы — Muon, остальное — Adam» можно заменить измерением на контрольной выборке, а перераспределение бюджета между блоками при фиксированной сумме даёт небольшое, но воспроизводимое улучшение. Оценка проведена на одной GPT-подобной модели и одном датасете, поэтому перенос результата на другие масштабы и данные остаётся открытым вопросом.