Scalable Second Order Optimization for Deep Learning
arxiv.org/abs/2002.09018

На обзор статьи меня навёл вот этот трэд в твиттере. Главный его момент был в том, что авторы DALLE-mini столкнулись с тем, что большие версии модели склонны внезапно расходиться посреди тренировки. Перепробовав кучу оптимизаторов и lr они обнаружили что некий Distributed Shampoo сходится гораздо быстрее и стабильнее Adam и Adafactor.

Главное отличие Dist Shampoo от Adam в том, что это метод оптимизации второго порядка, типа метода Ньютона, LBFGS или KFAC. Для апдейта Shampoo использует не только градиент G, но и матрицы ковариации градиента G G^T и G^T G, которые нужно ещё и возвести в степень 1/4 (возводить матрицы в степень это боль, а когда степень < 1 это боль в кубе).

Dist Shampoo добавляет кучу хаков, чтобы ускорить Shampoo. Первый хак: для матриц эмбеддингов, размер которых [vocab_size, hidden] используют только G G^T, тк вычислять произведение [vocab_size, hidden] @ [hidden, vocab_size] относительно дешево, а [hidden, vocab_size] @ [vocab_size, hidden] - безумно долго, тк vocab_size обычно раз в 50 больше hidden. Второй хак: в любом матричном произведении, если hidden очень большой, матрицу разбивают на N блоков размера hidden/N и вычисляют ковариации только этих блоков. Третий хак уже более хитрый: матрицы ковариации пересчитывают не на каждом шаге оптимизации, а только каждые ~100 шагов. Идея в том, что их рассчёт происходит на отдельной GPU/TPU параллельно обучению сети.

В оригинальной статье показали, что с помощью Dist Shampoo можно обучать трансформеры размером в 100M параметров за примерно в полтора раза меньшее время, чем Adam.

Подробнее про Distributed Shampoo:
1. WandB Report с DALLE mini
1. Имплементация на Jax