Adaptive Regularization through Coupled Kronecker Factoring

Adaptive regularization based optimization algorithms have been widely used in deep learning. While diagonal preconditioners such as Adam and Adagrad are scalable for large numbers of parameters, methods utilizing Kronecker factoring have recently found success. Notably, Shampoo, CASPR, and SOAP share identical axes preconditioners which are developed independently to optimize regret. However, there remain theoretical gaps regarding how to develop these preconditioner factors jointly and in an online fashion to minimize regret in online convex optimization. To address this, we first build a novel Loewner order constraint for approximate preconditioners, when followed, leads to a regret upper bound in online convex optimization. We show a correlation between LogDet matrix divergence of approximate preconditioner with full-matrix Adagrad and our regret upperbound. We then utilize both the Loewner order constraint and LogDet divergence to develop a Kronecker-structured preconditioner with factors jointly optimized to approximate full-matrix Adagrad. We show that a fully coupled Kronecker product is equivalent to Shampoo-KL, an existing algorithm, and propose a novel partially coupled version which uses the same number of matmuls as Shampoo, unlike Shampoo-KL. The Loewner order constraint and further analysis of the coupling between the axes preconditioners sets us up to show that fully coupled and partially coupled optimizers are optimal in online convex optimization. Our method improves generalization by up to 8% and accelerates training, improving time-to-target by up to 20%.