DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers
[Статья] [Код]

Оптимизатор с названием Шампунь (Shampoo, а есть еще и мыло 🧼) неплохо показал себя в ряде приложений - выдал лучшую сходимость по wall-clock time на AlgoPerf бенчмарке. Кроме того, обученные им модели устойчивы к квантизации.

Однако, шаг оптимизатора остается достаточно дорогостоящим из-за необходимости считать корни от матриц. В качестве обходного пути обновляют матрицы предобуславливания с какой-то периодичностью (раз в 100 шагов, скажем), но это сказывается негативно на сходимости.

В данной статье предлагают ускоренную и распределенную реализацию Шампуня - DASH (Distributed Accelerated SHampoo).