flâneur

KFAC explained

fdangel.com · 2,379 words · saved by 1 readers

This write-up accompanies the extension of Kronecker approximations to transpose convolution in BackPACK (Dangel, Kunstner, et al., 2020), which I have recently implemented. All approximations we will talk about tackle the Hessian of a neural network's loss. Let f θ 𝑓 𝜃 denote a neural network that maps a vector-valued input x 𝑥 to a vector-valued prediction f 𝑓 , which is then scored by a convex loss function ℓ(f,y)∈ℝ ℓ ( 𝑓 , 𝑦 ) ∈ 𝑅 , using the ground truth (label) y 𝑦 . The Hessian ∇ 2 θ ℓ ∇ 𝜃 2 ℓ collects the second-order derivatives of the loss w.r.t. the neural network's parameters θ 𝜃 and has elements More precisely, we will only be concerned with approximating certain blocks of this matrix. The block structure follows from the layer structure of our neural network. Assume that we have L 𝐿 layers with parameters θ (1) ,…, θ (L) 𝜃 ( 1 ) , … , 𝜃 ( 𝐿 ) such that The Hessian inherits this structure and consists of L 2 𝐿 2 blocks, where block ( l 1 ,

KFAC explained \(\def\mymacro{{\mathbf{\alpha,\beta,\gamma}}}\) \(\def\va{{\mathbf{a}}}\) \(\def\vb{{\mathbf{b}}}\) \(\def\vc{{\mathbf{c}}}\) \(\def\vd{{\mathbf{d}}}\) \(\def\ve{{\mathbf{e}}}\) \(\def\vf{{\mathbf{f}}}\) \(\def\vg{{\mathbf{g}}}\) \(\def\vh{{\mathbf{h}}}\) \(\def\vi{{\mathbf{i}}}\) \(\def\vj{{\mathbf{j}}}\) \(\def\vk{{\mathbf{k}}}\) \(\def\vl{{\mathbf{l}}}\) \(\def\vm{{\mathbf{m}}}\) \(\def\vn{{\mathbf{n}}}\) \(\def\vo{{\mathbf{o}}}\) \(\def\vp{{\mathbf{p}}}\) \(\def\vq{{\mathbf{q}}}\) \(\def\vr{{\mathbf{r}}}\) \(\def\vs{{\mathbf{s}}}\) \(\def\vt{{\mathbf{t}}}\) \(\def\vu{{\math

saved by

related reading