Эффективной стратегией при решении сложных задач является итеративное улучшение решения. Авторегрессионные языковые модели добиваются этого с помощью длинных цепочек рассуждений и для того, чтобы такой механизм эффективно работал, требуется дообучение на специальных данных или некая вариация RLя.
В качестве альтернативы можно держать выход некоторого фиксированного размера и итеративно его улучшать, обновляя раз за разом некоторое скрытое состояние.
HRM
В HRM 4 обучаемые компоненты:
📌 Входной эмбеддинг
f_I(x)📌 Низкоуровневая сеть
f_L(x + z_L + z_H), обновляющая состояние z_L
📌 Высокоуровневая сеть f_H(x + z_L + z_H), обновляющая состояние z_H📌Выходной эмбеддинг
f_O(z_H)Предсказание модель выдает не сразу, а через несколько прогонов. Сначала в цикле делается несколько прогонов через f_L, и один прогон через f_H. Прогоны делаются до определенной неким образом сходимости и последнее выходное состояние
z_H подается в выходную голову откуда получается предсказание для целевой задачи.Чтобы такую конструкцию можно было обучать эффективно по памяти, авторы делают предположение, что на последней итерации уравнение обновления скрытых состояний сошлось к фиксированной точке и можно делать приближение 1-го порядка, из которого следует, что достаточно считать градиенты только на последней итерации, а прошлые гонять с no_grad менеджером контекста. Тем самым не нужно хранить объемный граф вычислений.
В HRM максимально число итераций выставляется в 16, но бывает так, что можно сойтись гораздо раньше. И для этого обучают специальную голову поверх
z_H, которая говорит, пора ли останавливаться. За счет нескольких прогонов имеем эффективно большую глубину. Данный подход напоминает AlBERT, и Adaptive Computation Time из бородатых 🧔♂️ времен.
👶 TRM
В TRM предлагают переосмыслить архитектуру и процедуру обучения TRM. Вместо низкоуровневого и высокоуровневого состояния есть только одно скрытое состояние z и ответ y. Фактически z_L выполняет роль y, а z_H - z.
На практике сходимость и условия теоремы про фиксированную точку могут не выполняться и авторы переосмысляют алгоритм итераций так, чтобы можно было не исходить из этого предположения. Как и в HRM можно прогнать несколько итераций улучшения решения и сделать backprop на самой последней (но через все прогоны f_L). Данная модификация дает основной прирост метрик.
Далее авторы замечают, что не имеет смысла держать две модели (low-level и high-level) и одна модель работает не хуже.
Использование одного скрытого состояния z и выходного признака оптимально по качеству. Добавление большего числа скрытых состояний только ухудшает качество, якобы потому что это неестественно 🤷♂️ (более строго аргумента не приводится), а одно состояние на ответ тоже сильно хуже - мол, пытаемся в один вектор впихнуть и решение и скрытый признак. Логика такая, что модели нужно держать последнее актуальное решение и то, как она к нему пришла.
Хаки для улучшения качества
Делать модель больше не имеет смысла в виду ограниченности обучающей выборки. 2-слойная модель работает лучше всего, 4-слойная уже переобучается (регуляризации, правда, не особо шатали, по всей видимости).
Для малой длины контекста (рассматриваемые задачи имеют фиксированный размер входов) вместо Attention лучше себя показывает MLP-Mixer . На задачах побольше MLP-Mixer более склонен к переобучению, и потому берется стандартный трансформер.
Кроме того, для стабилизации обучения и борьбы с переобучением используют экспоненциальное скользящее среднее (EMA).
Количество рекурсий и отношение частоты обновление y и z настраивается для выжимания максимального качества.