DyKAF: Dynamical Kronecker Approximation of the Fisher Information Matrix for Gradient Preconditioning

TL;DR

DyKAF introduces a novel, efficient method for approximating the Fisher Information Matrix using projector-splitting integrators, leading to improved optimization in large language model training.

Contribution

The paper presents DyKAF, a new optimizer that enhances Fisher matrix approximation with a dynamical, Kronecker-based approach, outperforming existing methods.

Findings

DyKAF improves Fisher matrix approximation quality.

DyKAF outperforms existing optimizers in language model tasks.

The method is effective in large-scale pre-training and fine-tuning.

Abstract

Recently, optimizers that explicitly treat weights as matrices, rather than flattened vectors, have demonstrated their effectiveness. This perspective naturally leads to structured approximations of the Fisher matrix as preconditioners, where the matrix view induces a Kronecker-factorized form that enables memory-efficient representation. However, constructing such approximations both efficiently and accurately remains an open challenge, since obtaining the optimal factorization is resource-intensive and practical methods therefore rely on heuristic design choices. In this work, we introduce a novel approach that leverages projector-splitting integrators to construct effective preconditioners. Our optimizer, DyKAF (Dynamical Kronecker Approximation of the Fisher Matrix), consistently improves the Fisher matrix approximation quality. Experiments on large language model pre-training and fine-tuning demonstrate that DyKAF outperforms existing optimizers across a range of evaluation metrics.