Всем привет! В этом посте я расскажу о проекте, который выполнил совместно с командой Google Brain во время исследовательской стажировки в Цюрихе. Мы работали над моделью обработки естественного языка, которая решает задачи на общее понимание текста (задачи из набора GLUE: General Language Understanding Evaluation). BERT-подобные модели мы комбинировали с помощью маршрутизирующих сетей и добились того, что при увеличении мощности скорость вывода почти не изменилась. Финальная модель объединяет 10 BERTlarge моделей и имеет более 3,4 миллиарда параметров. Подробности под катом!  Меня зовут Никита Сазанович, я учусь на 2-м курсе магистратуры «Программирование и анализ данных» в Питерской Вышке. В январе я поехал на 17-недельную стажировку в Google Brain (правда, физически находился в Цюрихе только 11 из них: из-за карантина последние 6 недель пришлось работать удаленно). Мой проект был связан с обработкой естестественного языка с использованием маршрутизирующих сетей — а это новое понятие в машинном обучении, с которым мы далее и разберемся. Задачи обработки естественного языка (Natural Language Processing, NLP) включают в себя машинный перевод, распознавание текста, анализ тональности текстов и др. В этой статье мы поговорим про задачи из набора GLUE. Они были разработаны экспертами машинного обучения, чтобы оценивать, насколько модель способна к общему пониманию текста (откуда и название набора: General Language Understanding Evaluation).  Всего в наборе GLUE 9 задач. Все они получают на вход одно или пару предложений и либо разбивают их на два или три класса, либо проводят регрессию. По цели они разделены на три категории: задачи над одним предложением, задачи сходства и перефразирования и задачи логического вывода. Рассмотрим одну из них: Quora Question Pairs, или QQP. Как ясно из названия, нужно будет что-то сказать про пару вопросов из сайта Quora. Вопросы необходимо сравнить и сообщить, семантически эквивалентны ли они / является ли один дубликатом другого. Сразу приведу два примера: 1) Вопросы "What are natural numbers?" и "What is a least natural number?" не являются дубликатами. 2) А вопросы "How do you start a bakery?" и "How can one start a bakery business?" семантически эквивалентны. То есть в этой задаче по тексту двух вопросов просят провести бинарную классификацию: разбить их на два класса. Задачу классификации в машинном обучении решают уже давно и научились это делать хорошо, но не в тех случаях, когда дело касается обработки естественного языка. Можно решить эту проблему, если привести текст к формату, на котором осуществлять классификацию более удобно. Например, преобразовать его в единый числовой вектор, который и послужит входом для классификации! Такой вектор называют представлением языка, а получением таких векторов занимается область обучения представлений языка, или language representation learning. Итак, для того чтобы оценить, насколько модель способна к общему пониманию текста, ей предлагают решить несколько задач из набора GLUE, а для того чтобы она могла решать эти задачи, нужно подать ей данные в удобном формате — в формате представления языка (то есть вектора). Далее мы поговорим об одном из современных подходов к получению представлений языка. Обычно к обработке текста подходят с помощью токенизации, то есть разбиения предложения на токены из заданного словаря. Под токеном понимается определенная последовательность символов. Токены и словари строятся различными способами, мы будет говорить о подходе WordPiece. При таком подходе слово, которое не встречается в словаре целиком, будет разбито на части так, чтобы каждая часть была токеном словаря WordPiece. Пример такой токенизации: 1) Входное предложение: "Jet makers feud over seat width with big orders at stake". 2) Разбиение на токены (где символ _ приписывается к токенам, с которых слова начинаются): "_J et _makers _fe ud _over _seat _width _with _big _orders _at _stake". Здесь можно видеть, что слов "Jet" и "feud" не оказалось целиком в словаре, поэтому они были разбиты на части, присутствующие в словаре. В словаре WordPiece содержится несколько десятков тысяч токенов. Большинство слов, которые встречаются наиболее часто, он содержит целиком. Итак, мы разбили текст на токены! Для получения представления текста нам осталось как-то отобразить каждый из этих токенов в вектор, а затем скомбинировать вектора в один. Первый шаг большинство моделей осуществляют, храня словарь векторов представлений для каждого токена в отдельности, без контекста. Но многие исследования проводят для ответа на вопрос: как организовать обучение, чтобы в итоге представления для токенов учитывали контекст текста, в котором они находятся? В 2017 году был представлен механизм внимания, который умеет учитывать контекст. Еще один плюс механизма в том, что зависимости в обученной модели легко интерпретировать для каждого входного текста. Рассмотрим пример на картинке снизу:  Если мы посмотрим на слово «making» в предложении «It is in this spirit that a majority of American governments have passed new laws since 2009 making the registration or voting process more difficult», мы можем обратить внимание также на слова «more difficult», «laws», «2009»: они помогают понять смысл. В предложении говорится о том, как с 2009 года правительства США принимали законы, которые делают процесс голосования сложнее. В таком контексте представление слова «making» в механизме внимания формируется в основном из представлений предыдущего слоя для слов «making», «more» и «difficult». Математически механизм внимания представляет собой блок модели, который отображает последовательность векторов токенов. При этом блок можно устроить таким образом, чтобы выходное представление каждого токена формировалось как взвешенная сумма всех представлений токенов с предыдущего слоя. В каком-то смысле у каждого слова есть ограниченное «внимание», которое оно разделяет между всеми словами предыдущего слоя. В той же статье авторы показали модель, которая состояла из нескольких последовательных блоков с механизмом внимания, и назвали ее Трансформер. Позже был представлен способ предобучения представлений BERT, который использует модель Трансформера, а также описывает датасет и функцию обучения. Под BERT-ом также часто имеют в виду предобученные модели, которые опубликовали авторы метода, и которые можно очень легко и быстро дообучать для вашей задачи. Вернемся к вопросу представления текста. Допустим, после каких-то преобразований мы получили последовательность векторов , , ...,  для токенов текста , , ..., . Как теперь проводить классификацию или регрессию на всем тексте? В BERT модели предлагают элегантное решение: к каждой последовательности токенов добавить токен  и получить по нему представление  описанным выше способом. Благодаря специальному устройству функции обучения оно будет представлять информацию, которая содержится во всем тексте. Для более подробного устройства можно обратиться к самой статье о BERT. Ниже я расскажу о том, как эти BERT-подобные модели можно комбинировать с помощью нового для машинного обучения понятия маршрутизирующих сетей при решении набора задач GLUE. Под BERT-подобными моделями мы будет подразумевать BERTы, которые имеют разные гиперпараметры и/или обучены на отличном от оригинального датасете (примеры таких моделей).  Вот простая аналогия для понимания таких сетей: когда вы решаете какую-то проблему или задачу, вы можете захотеть обратиться к какому-то количеству «экспертов». Ими могут быть ваши родственники, коллеги или друзья. В зависимости от типа задачи вы будете обращаться к разным экспертам: для задач с работы, скорее всего, — к коллегам, а для вопросов, связанных с хобби, — к друзьям. При этом набор экспертов остается постоянным. Суть маршрутизирующих сетей заключается в том, что обучение происходит на нескольких задачах одновременно, при этом для каждой задачи выбирается набор модулей/маршрут в сети модулей (наборе экспертов). В маршрутизаторах физических сетей маршрут зависит от оценки эффективности маршрутизации через тот или иной участок. В маршрутизирующих сетях машинного обучения он тоже не фиксирован и изменяется в процессе тренировки. Таким образом, одновременно происходит обучение как составляющих модели, так и маршрутизатора.  Формулировать и обучать процесс маршрутизации можно по-разному. К примеру, в статье по ссылке для каждой возможности маршрутизации задачи через модуль модели использовалось распределение Бернулли, которое обучалось с помощью Gumbel Softmax. Ниже вы можете видеть пример такой модели с двумя задачами. Цветами отмечено, какое подмножество модулей используется для каждой из них: красным обозначены модули, которые вычисляются на маршруте для первой задачи, синим — для второй, а фиолетовым — для обеих. Если в одном слое таких модулей для одной задачи несколько, нужна агрегация результатов, например, усреднение. Наличие общих модулей и позволяет задачам совместно обучать модули, которые выгодно разделять.  Как видно на картинке, модули этой сети состоят из сверточных слоев (Conv 3x3, Conv 5x5, Conv 7x7). Значит, задачи относились к области компьютерного зрения. А что если в качестве таких модулей мы будем использовать BERT-подобные модели и попытаемся искать маршруты для задач из набора GLUE? Это и стало идеей для моего проекта. На вход модель получает токенизированный текст (x) и числовой идентификатор задачи (t), который превращается в вектор представления задачи. Используя этот вектор, маршрутизатор выбирает набор модулей/экспертов для задачи, которым подается на вход токенизированный текст x. После обработки экспертами происходит агрегация результатов и получается представление текста. Затем представление направляется в небольшую сеть (отдельную для каждой задачи t), которая преобразует представление к ответу задачи (регрессии или классификации). Стоит также заметить, что для задач GLUE эксперты инициализируются предобученными моделями.  Представления задач моделируются как вектора размерности M (в экспериментах M было равно 10):   Маршрутизатор, используя этот вектор, вычисляет логиты коэффициентов каждого эксперта для этой задачи, после чего остаются K наибольших. Они нормируются функцией softmax:   Сами представления языка формируются как взвешенная сумма результатов экспертов, причем эксперту не нужно вычислять результат, если коэффициент равен нулю:   Изначально мы предпочитаем опрашивать малое число экспертов, чтобы объем вычислений был сравним с тем, который требуется при использовании одного эксперта. Как видно из устройства модели, если наш маршрутизатор при любом количестве экспертов всегда выбирает два пути, то объем вычислений для одного запроса будет лишь в два раза превышать объем для одного эксперта, а скорость вывода после обучения не будет превышать времени вывода двух экспертов на одном устройстве. Это и отвечает на вопрос об устройстве модели обучения, поставленный в самом начале статьи. Далее под спойлером для любителей описаны технические подробности настройки этой модели. Для полноты картины расскажу про нее в нескольких словах. Лучший результат достигается, если использовать 10 BERT large экспертов, а не трех или одного (из-за затрат времени и ресурсов мы провели эксперименты только для 1, 3 и 10 экспертов). Оптимальное количество путей — два. Оптимизировать выгодно стохастическим градиентным спуском с импульсом. При такой настройке модель имеет 3,4 миллиарда параметров и обучать ее пришлось на 24 тензорных процессорах Google. Процесс тренировки проходит параллельно на батчах из всех задач GLUE и требует какой-то агрегации между задачами. Один из вариантов — линейная комбинация потерь  каждой из задач:     В первых трех строках можно видеть разные способы задания коэффициента, который может включать количество задач или размер данных для задачи. Ниже представлены два способа, которые имеют другую форму (статья про последний способ). Лучшим из пяти способов оказался третий вид, но если рассмотреть общую форму с нормированием до единицы и учетом размера датасета:    Важным элементов при масштабировании модели является метод оптимизации, так как он увеличивает размер необходимой памяти, тем самым уменьшая мощность модели при фиксированном бюджете. Я проверил 6 оптимизаторов: Adam, Adam с weight decay (угасание весов к нулю), LAMB, LAMB с weight decay, стохастический градиентный спуск (SGD) и стохастический градиентный спуск с импульсом (SGDM). SGD практически не расходует дополнительной памяти, SGDM запоминает импульсы для каждой переменной, поэтому расходует x2 памяти, а остальные рассмотренные оптимизаторы расходуют x3 памяти. Результаты оптимизации маршрутизируемого BERT-а с одним BERTlarge экспертом можно видеть в таблице. SGDM достигает лучшего результата в исследуемом пространстве, при этом используя лишь в два раза больше памяти. Я использую его в финальной модели.  Рассматривая зависимость качества от числа BERTlarge экспертов и путей, я заметил, что качество улучшается при увеличении числа экспертов, а при увеличении количества путей — не всегда. Лучший результат дало использование 10 экспертов и 2 путей.  В такой модели 10 x 340 миллионов = 3,4 миллиарда параметров, а ее обучение происходило на 24 TPUv2 (тензорных процессорах Google) с параллелизмом по данных равным 4 и параллелизмом модели равным 6. Параллелизм модели равен 6, так как 5 процессоров держат по 2 BERTlarge модели, а еще один отвечает за все сети задач после агрегации. Для сравнения результатов в наборе GLUE был предложен способ комплексной оценки, который совмещает результаты на каждой из задач набора. Эта оценка лежит в диапазоне от 0 (все неверно) до 100 (идеальные ответы). Я сравнил маршрутизируемый BERT с опубликованными результатами мультизадачного обучения BAM! (оно основано на дистилляции) и настройкой одного BERTа под каждую задачу. Маршрутизируемый BERT достиг лучшей оценки общего понимания языка. Интересны также различия между результатами на задачах. Например, выгоду при тренировке с маршрутизируемыми сетями получает задача RTE, которая имеет наименьший тренировочный набор (2,5 тысячи примеров). Сравните ее оценки в 1-й и 3-й строках. Эта выгода на задаче достигается засчет использования дополнительных знаний из других задач, тренировочный набор которых больше по размеру.  Еще один результат этой модели — возможность исследовать отношения между задачами. Ниже представлен результат сокращения размерности векторов для представлений задач. Вектора, которые изначально находятся в 10-мерном пространстве (M=10), сокращены до векторов размера 2, т. е. точек. Это сделано при помощи метода t-SNE, который очень популярен для таких визуализаций и про который можно почитать в другом блоге на Хабре. Здесь каждая из задач располагается на плоскости, а расстояние между ними позволяет судить о том, насколько они близки с точки зрения маршрутизатора. Как и ожидалось, mnli-m и mnli-mm — задачи с одним и тем же тренировочным набором — располагаются близко. Так же близко расположены sst-2 и qqp, cola и rte — пары, сходство которых отмечали и другие исследователи.  В результате мы разобрались, как можно использовать маршрутизирующие сети для обучения представлений языка. С их помощью мы улучшили результат работы исходной модели с точки зрения оценки GLUE на 1,75 пункта и в качестве бонуса смогли визуализировать отношения между задачами набора. Важное свойство такого подхода — сохранение скорости вывода в конечной модели. Было бы интересно продолжить работу над проектом, использовав разнородных экспертов с одинаковым интерфейсом (таких как ALBERT и RoBERTa), а также добавить к набору задачи предобучения представлений (таких как MLM BERT-а) для тренировки сети с нуля без использования предобученных моделей. Код этой модели остался в Google, но могу сказать, что написан он был на любимом в компании TensorFlow. Магистр ВШЭ СПб