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

Исследователи Amazon Web Services представили на конференции ICML метод ControlG. Он помогает обучать нейросеть сразу нескольким задачам, которые мешают друг другу. Идею авторы взяли из промышленной автоматики: очерёдность задач определяет ПИД-регулятор, тот же тип контроллера, что держит скорость в круиз-контроле автомобиля или температуру воды в дорогой кофемашине. Основную часть работы сделал Karish Grover, аспирант Carnegie Mellon University, во время стажировки в Amazon. Код выложен в открытый доступ.

Проверяли метод на графовом обучении без учителя: модель учится понимать графы (узлы и связи между ними) без размеченных данных. Для этого обычно используют сразу несколько целей обучения. Модель предсказывает пропущенные связи, восстанавливает скрытые признаки узлов или учится находить общее между узлом и его соседями. Ни одна цель не выигрывает на всех данных, поэтому их комбинируют. Стандартный способ — на каждом шаге обучения смешивать поправки ко всем целям в одно обновление. Авторы видят в этом три проблемы. Цели тянут параметры в разные стороны, и успехи в одной портят другую. Полезная в начале цель к концу обучения становится лишней, но веса за этим не успевают. А адаптивные схемы взвешивания иногда сводят вес какой-то цели почти к нулю, и она фактически выпадает из обучения.

Авторы предлагают другую постановку. Вопрос не в том, как смешать цели прямо сейчас, а в том, какой из них отдать следующую порцию вычислений. Даже случайная очерёдность неожиданно оказалась сильной: в кластеризации узлов её средний ранг 5,0, тогда как у сложных методов AutoSSL, PCGrad, CAGrad и ParetoGNN он от 7,3 до 10,0. ControlG распределяет вычисления осмысленно, на трёх уровнях. Раз в какое-то время он оценивает трудность каждой цели. Учитываются два показателя: насколько сильно противоречат друг другу обновления соседних узлов и насколько цель конфликтует с остальными. Раз в эпоху он составляет план, какую долю вычислений отдать каждой цели, и в приоритете оказываются отстающие. В промежутках ПИД-регулятор следит за выполнением плана и не даёт распределению раскачиваться.

Метод испытали на девяти графовых наборах данных, самый крупный из которых — граф цитирований научных статей на 169 тыс. узлов, в трёх задачах: классификации узлов, предсказании связей и кластеризации. Средние ранги ControlG — 1,4, 1,9 и 1,8, это лучше всех конкурентов. На самом крупном графе точность классификации составила 72,86% против 71,62% у ближайшего многозадачного конкурента CAGrad. Дополнительная нагрузка — 16–31 мс на шаг. Это больше, чем у простого планирования (8–15 мс), но намного меньше, чем у AutoSSL (125–414 мс) и ParetoGNN (35–764 мс). Если убрать планировщик, качество падает до 3,4%. Если заменить регулятор случайной выборкой по плану, потери составляют 1–2%.

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