Official PyTorch implementation of Mei, Jonathan, Alexander Moreno and Luke Walters. "KrADagrad: Kronecker Approximation-Domination Gradient Preconditioned Stochastic Optimization." Uncertainty in Artificial Intelligence (2023).
- Main code lives in
kradagrad.py
(KrADagrad* in the paper) andkradagradmm.py
(KrADagrad in the paper) - We include
shampoo
inthird_party/
with minor modifications for convenience - An implementation of resnet also lives in
third_party/
for convenience - Methods for taking p-th roots (and other useful utilities) for symmetric positive semi-definite matrices live in
positive_matrix_functions.py
- Experiments (training notebooks+scripts) are found in
experiments/
Set up environment using conda:
conda env create -f krad.yml