An Optimizer for Conformer Based on Conjugate Gradient Method

Optimization is a fundamental problem in AI research. Widely used first-order optimization methods like SGD and Adam have drawbacks, such as difficulty in hyper-parameter selection, poor performance on “ill-conditioned” loss landscapes, and bottle-neck of efficiency improvement. This has led to interest in “second-order” optimization methods, which, however, face high space and time complexity due to the storage and inversion of second-order matrices. Designing practical second-order methods to overcome these complexities is crucial. The Hessian-free method based on the conjugate gradient method (HFCG) addresses these issues by approximating the “second-order” updates iteratively. This paper extends HFCG to the Conformer model, implementing directional derivative calculations for multi-head self-attention, convolution, LayerNorm, and BatchNorm modules for the first time. A two-stage distributed training procedure (TDTP) is introduced to enhance optimization performance and reduce training time, including an improved update strategy and a new CG batch extraction method. TDTP was tested on the Penn Treebank and AMI datasets for language modelling tasks. Ablation experiments demonstrated the efficiency of HFCG and the new TDTP strate-gies. Additionally, the impact of the choice of <tex xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">$\mathbf{H}_{L}$</tex> in the generalized Gauss-Newton matrix on performance was explored and the hypothesis that HFCG might reduce performance differences caused by normalization was proposed.