DL in NLP

12 Jan 2021 20:33

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
arxiv.org/abs/2101.03961
Fedus et al. [Google Brain]

Twitter thread: тык
Код: тык

TL;DR Допилили напильником mixture of experts, чтобы натренировать модель с 1.6 триллионами параметров.

Рисёч в области больших неросетей уже больше рисёч в области high-performance computing, чем в области нейросетей: вам нужно, чтобы ваши 100500 GPU друг с другом эффективно коммуницировали и чтобы каждая из них использовалась в полную силу. Одним из путей решения этих проблем (и создания новых) является подход Mixture of Experts (MoE), который уже позволил натренировать трансформер с 1 триллионом парметров. Основная идея: в каждом (втором) слое транфмораера FFN слой заменяется на MoE-слой, который представляет собой N (~1024) очень широких полносвязных сеток-экспертов (у них разные параметры). Каждая из сеток может лежать на своей GPU и входящая последовательность токенов в начале попадает в сетку-роутер, которая с помощью софтмакса определяет какой токен отдать каким k экспертам.

Авторы Switch Transformers дорабатывают этот подход:
1. Используют k=1 - то есть роутер отдаёт каждый токен только в одного эксперта. В оригинальной MoE-статье говорилось, что это не должно работать, тк градиенты становятся тривиальными. Но у Switch Transformers заработало 💁‍♂️.
1. Перед расчётом софтмакса в роутере конвертируют тензоры в fp32. Вся остальная сетка работает в fp16.
1. Чуть-чуть поменяли вид регуляризации, которая заставляет роутер использовать всех экспертов и не концентрироваться только на нескольких из них.
1. Инициализируют модель в sqrt(10) раз меньшими числами.
1. Если в эксперта приходит слишком много токенов, часть этих токенов пропускает эксперта (по сути делается skip-connection).

Результаты:
1. Обучается с меньшим количеством танцев с бубном, чем MoE - позволяют натренировать модель с 1.6 триллионами параметров без проблем с нестабильностью
1. На 600% раз быстрее, чем T5 с эквивалентным FLOPS/token и на 20% быстрее, чем MoE.
1. На языковом моделировании и на NLU-датасетах типа SuperGLUE заметно улучшает результаты T5 (размер моделей подобран так, что на их предобучение уходит одинаковое время).
1. Авторы предполагают, что даже относительно небольшие модели могут работать лучше и обучаться быстрее, если использовать Switch Transformers.

Статья очень большая и читается сложновато. Если решите вникнуть поглубже, рекомендую обратить внимание на Appendix F с псевдо-TF-кодом.

Ещё вот неплохая имплементация MoE на PyTorch: тык