Code to reproduce the results to the ICLR 2026 paper "On the Convergence Behavior of Preconditioned Gradient Descent Toward the Rich Learning Regime" (https://arxiv.org/pdf/2601.03162)

Figure 1:experiments/grokking/mnist-grokking-best.ipynb
Figure 3:experiments/gauss_newton/ex_1.ipynb
Figure 4:experiments/pinns/make_plots_new.ipynb
Figure 5:experiments/grokking/modulo_example.ipynb
Figure 6:experiments/grokking/grokking_polynomial_good.ipynb
Figure 7:experiments/grokking/mnist-grokking-best.ipynb
Figure 8:experiments/grokking/modular-plotter.ipynb
Figure 9: experiments/gauss_newton/ex_2d_discont.ipynb
Figure 10:experiments/grokking/mnist-grokking-best.ipynb
Figure 11:experiments/grokking/mnist-plotter-xentropy.ipynb
Figure 12:experiments/grokking/modular-plotter.ipynb
Figure 13: Adjust seed parameters from above.