2022-01-17 22:40:20
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
1.2K viewsVlad Lialin, 19:40
D
