ASGO: Adaptive Structured Gradient Optimization
This work addresses the inefficiency of current optimizers like Adam in leveraging structural properties of gradients and Hessians in deep learning, offering a domain-specific improvement for training neural networks.
The paper tackles the problem of training deep neural networks by proposing ASGO, an optimization algorithm that uses adaptive structured gradients to achieve superior convergence rates compared to existing methods, with empirical verification on language model tasks.
Training deep neural networks is a structured optimization problem, because the parameters are naturally represented by matrices and tensors rather than by vectors. Under this structural representation, it has been widely observed that gradients are low-rank and Hessians are approximately block diagonal. These structured properties are crucial for designing efficient optimization algorithms, but are not utilized by many current popular optimizers like Adam. In this paper, we present a novel optimization algorithm ASGO that capitalizes on these properties by employing a preconditioner that is adaptively updated using structured gradients. By a fine-grained theoretical analysis, ASGO is proven to achieve superior convergence rates compared to existing structured gradient methods. Based on this convergence theory, we further demonstrate that ASGO can benefit from low-rank gradients and block diagonal Hessians. We also discuss practical modifications of ASGO and empirically verify ASGO's effectiveness on language model tasks. Code is available at https://github.com/infinity-stars/ASGO.