Training Invertible Linear Layers through Rank-One Perturbations
This work addresses a specific problem in machine learning for researchers and practitioners using invertible layers, offering an incremental improvement over existing reparameterization or manifold optimization methods.
The paper tackles the challenge of maintaining invertibility in linear layers during neural network training by introducing a method that trains rank-one perturbations and adds them to weight matrices infrequently, enabling efficient tracking of inverses and determinants without explicit computation. The approach improves mixing and mode separation in normalizing flows, as demonstrated experimentally.
Many types of neural network layers rely on matrix properties such as invertibility or orthogonality. Retaining such properties during optimization with gradient-based stochastic optimizers is a challenging task, which is usually addressed by either reparameterization of the affected parameters or by directly optimizing on the manifold. This work presents a novel approach for training invertible linear layers. In lieu of directly optimizing the network parameters, we train rank-one perturbations and add them to the actual weight matrices infrequently. This P$^{4}$Inv update allows keeping track of inverses and determinants without ever explicitly computing them. We show how such invertible blocks improve the mixing and thus the mode separation of the resulting normalizing flows. Furthermore, we outline how the P$^4$ concept can be utilized to retain properties other than invertibility.