Knowledge Accumulator

Energy-Based Transformers are Scalable Learners and Thinkers [2025]

Начнём с базы. Обычно в ML преобладают модели типа p(y | x) - то есть, на вход подаются фичи, а на выходе распределение на ответ. У этого есть альтернатива - мы можем иметь модель вида E(x, y) - так называемую энергетическую функцию.

E(x, y) показывает "правильность" пары этих двух объектов - если энергия низкая, то пара подходит друг к другу, а если энергия высокая, то не подходит. Теоретически, эти две схемы эквивалентны - можно считать, что обычный классификатор p(y | x) выдаёт энергию для всевозможных y.

Но на практике E(x, y) обладает большей мощностью при том же количестве параметров. Всё дело в том, что один подсчёт p(y | x) эквивалентен множеству применений E(x, y) для каждого ответа. E(x, y) должна лишь оценить качество ответа, тогда как p(y|x) должна решать задачу, последнее фундаментально сложнее.

Рекомендательные системы, кстати, прекрасно впитали эту идею. На последних стадиях ранжирования мы применяем именно модели вида F(user, item), а не p(item | user). Тут же становится очевидной их главная проблема - для выбора ответа функцию надо применять на каждом y,

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

Введение можно считать оконченным...

Energy-Based Transformers делает ровно последнее. Мы подаём на вход последовательность токенов, но вместо того, чтобы выдавать распределение на следующий токен на выходе, мы подаём это распределение на вход, и предсказываем энергию этой пары. Во время инференса решаем задачу оптимизации по этому распределению.

Танцы с бубном начинаются именно при обучении энергетической функции - у неё ведь на входе теперь ещё и распределение на все классы. Как и в рекомендациях, в таких моделях можно применять тот самый Constrastive loss, про который я говорил в посте про LogQ correction. Но учить на всём декартовом произведении X и Y это кошмарно дорого, а вообще-то нам нужно выдавать энергию для распределения на Y.

В статье к обучению E(x, p(y)) подходят в лоб. Процедура инференса - оптимизация вектора p(y) градиентным спуском по E(x, p(y)) - это, на секундочку, полностью дифференцируемая операция. Делаем так. Берём случайное p(y). Вычитаем из него N раз градиент E(x, p(y)) по p(y), и потом считаем итоговый лосс как близость финального p(y) и правильного p(y).

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

Перейдём к результатам.

Самое главное - точность на валидации скейлится значительно быстрее с увеличением объёма тренировочных данных и уже обгоняет её на 3 миллиардах токенов. С другой стороны, нужно почти на порядок больше флопсов, чтобы получить то же качество. Учитывая то, что компьют можно удешевлять бесконечно, тогда в как разумных данных мы ограничены, то скейлинг с точки зрения данных является более важным показателем.

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

@knowledge_accumulator

1👍13🔥4❤2

2.97K viewsedited 10:51