summaries: This paper proposes a model for continual learning. Specifically, it intends to merge parameter regularization methods and gradient projection methods. The regularization is achieved by taking a bayesian perspective over the distribution over parameters, that is approximated by a Laplace distribution in a variational fashion, with online updates. The gradient projection is addressed by trust-region optimization, resulting in an update rule that uses the precision matrix to precondition the gradient of the current task objective, plus a term that keeps the parameters close to the previous mean. Experiments are carried out for sequential problems on a stimulus-response dataset and on stroke MNIST.

summaries: The paper addresses the problem of catastrophic forgetting during sequential task learning with recurrent neural networks, a notoriously difficult continual learning setting.   The main theoretical contribution is a formalism, Natural Continual Learning (NCL), which generalizes several related continual learning approaches in the regularization-based family, including several recent SOTA techniques. This formalism applies to the probabilistic learning setting using neural networks, so it is not restricted to supervised learning or recurrent models.  Experimental evidence is provided to support the claim that NCL outperforms previous approaches.  Finally, interesting analogies are made to models of biological learning, which are also supported by accompanying experiments; similar analysis tools are applied to practical instantiations of the NCL framework in controlled experimental settings.

summaries: This paper proposes a novel regularization based continual learning approach, NCL, which is mainly used for recurrent neural network (RNN). Since NCL is based on Bayesian continual learning, NCL utilizes the approximate posterior of previous tasks as a prior for current task, and the Laplace approximation is used to compute the approximate posterior. Through simple toy analyses on convex and non-convex loss landscapes for NCL and other baseline methods, authors show that previous methods have a problem on the trajectories of convergence when learning a new task. Based on the results on simple toy analyses, authors proposed a novel trust region based regularization method that modulate the gradient by considering the curvature of prior distribution via its precision matrix. The proposed technique is highly similar to projection based methods, but its origin is quite different from previous methods. Experiment results show that NCL outperforms other baselines in stimulus-response and MNIST tasks.

summaries: The authors propose a method called Natural Continual Learning (NCL). NCL unifies the concepts of weight regularization and gradient projection into non-interfering subspaces. Gradient projection is implemented as a trust region optimization using the Fisher information matrix. The authors also introduce a Kronecker-factor approximation to improve the scalability properties of their method. The comparison against other continual learning approaches is conducted using recurrent neural network (RNN) architectures.

summaries: The authors introduce an extension to Bayesian continual learning whereby updates in successive tasks are constrained to be close to a “trust region” of the parameters found in preceding tasks. Exploiting the Bayesian nature of the model, the trust region in question can be defined in terms of the prior precision matrix. The authors show that the method, when applied to RNNs, is effective against baselines in a) a suite of neuroscience related tasks, and b) stroke MNIST datasets.

abstractText: Biological agents are known to learn many different tasks over the course of their lives, and to be able to revisit previous tasks and behaviors with little to no loss in performance. In contrast, artificial agents are prone to ‘catastrophic forgetting’ whereby performance on previous tasks deteriorates rapidly as new ones are acquired. This shortcoming has recently been addressed using methods that encourage parameters to stay close to those used for previous tasks. This can be done by (i) using specific parameter regularizers that map out suitable destinations in parameter space, or (ii) guiding the optimization journey by projecting gradients into subspaces that do not interfere with previous tasks. However, these methods often exhibit subpar performance in both feedforward and recurrent neural networks, with recurrent networks being of interest to the study of neural dynamics supporting biological continual learning. In this work, we propose Natural Continual Learning (NCL), a new method that unifies weight regularization and projected gradient descent. NCL uses Bayesian weight regularization to encourage good performance on all tasks at convergence and combines this with gradient projection using the prior precision, which prevents catastrophic forgetting during optimization. Our method outperforms both standard weight regularization techniques and projection based approaches when applied to continual learning problems in feedforward and recurrent networks. Finally, the trained networks evolve task-specific dynamics that are strongly preserved as new tasks are learned, similar to experimental findings in biological circuits.

authors:
R. Aljundi
F. Babiloni
M. Elhoseiny
M. Rohrbach
T. Tuytelaars

title: Memory aware synapses: Learning what (not) to forget

venue: In Proceedings of the European Conference on Computer Vision (ECCV),

title: Natural gradient works efficiently in learning

venue: Neural Computation,

authors:
K.C. Ames
M.M. Churchland

title: Motor cortex signals for each arm are mixed across hemispheres and neurons yet partitioned within the population response. eLife, 8:e46159

authors:
A. Bernacchia
M. Lengyel
G. Hennequin

title: Exact natural gradient in deep linear networks and its application to the nonlinear case

venue: Advances in Neural Information Processing Systems,

authors:
A. Chaudhry
P.K. Dokania
T. Ajanthan
P.H. Torr

title: Riemannian walk for incremental learning: Understanding forgetting and intransigence

venue: arXiv preprint arXiv:1801.10112

authors:
C.A. Chestek
A.P. Batista
G. Santhanam
M.Y. Byron
A. Afshar
J.P. Cunningham
V. Gilja
S.I. Ryu
M.M. Churchland
K.V. Shenoy

title: Single-neuron stability during repeated reaching in macaque premotor cortex

venue: Journal of Neuroscience,

authors:
C. Clopath
T. Bonhoeffer
M. Hübener
T. Rose

title: Variance and invariance of neuronal long-term representations

venue: Philosophical Transactions of the Royal Society B: Biological Sciences,

authors:
Y. Cong
M. Zhao
J. Li
S. Wang
L. Carin

title: GAN memory with no forgetting

title: Incremental sequence learning. arXiv preprint arXiv:1611.03068

authors:
M. De Lange
T. Tuytelaars

title: Continual prototype evolution: Learning online from non-stationary data streams

venue: arXiv preprint arXiv:2009.00919

authors:
A.K. Dhawale
R. Poddar
S.B. Wolff
V.A. Normand
E. Kopelowitz
B.P. Ölveczky

title: Automated long-term recording and analysis of neural activity in behaving animals. eLife, 6:e27702

authors:
L. Duncker
L. Driscoll
K.V. Shenoy
M. Sahani
D. Sussillo

title: Organizing recurrent network dynamics by task-computation to enable continual learning

authors:
B. Ehret
C. Henning
M.R. Cervera
A. Meulemans
J. von Oswald
B.F. Grewe

title: Continual learning in recurrent neural networks with hypernetworks

venue: arXiv preprint arXiv:2006.12109

authors:
S.W. Failor
M. Carandini
K.D. Harris

title: Learning orthogonalizes visual cortical population

venue: codes. bioRxiv

authors:
J.A. Gallego
M.G. Perich
R.H. Chowdhury
S.A. Solla
L.E. Miller

title: Longterm stability of cortical population dynamics underlying consistent behavior

venue: Nature Neuroscience,

authors:
K. Ganguly
J.M. Carmena

title: Emergence of a stable cortical map for neuroprosthetic control

venue: PLoS Biology,

title: On quadratic penalties in elastic weight consolidation

venue: arXiv preprint arXiv:1712.03847

authors:
K. Jensen
Kao
T.-C
M. Tripodi
G. Hennequin

title: Manifold GPLVMs for discovering non-Euclidean latent structure in neural data

authors:
K.T. Jensen
N. Kadmon Harpaz
A.K. Dhawale
S.B.E. Wolff
B.P. Ölveczky

title: Long-term stability of neural activity in the motor system. bioRxiv

authors:
K.T. Jensen
Kao
T.-C
J.T. Stone
G. Hennequin

title: Scalable Bayesian GPFA with automatic relevance determination and discrete noise models. bioRxiv

authors:
K.A. Katlowitz
M.A. Picardo
M.A. Long

title: Stable sequential activity underlying the maintenance of a precisely executed skilled behavior

authors:
M.T. Kaufman
M.M. Churchland
S.I. Ryu
K.V. Shenoy

title: Cortical activity in the null space: permitting preparation without movement

authors:
D.P. Kingma
J. Ba

title: Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980

authors:
J. Kirkpatrick
R. Pascanu
N. Rabinowitz
J. Veness
G. Desjardins
A.A. Rusu
K. Milan
J. Quan
T. Ramalho
A Grabska-Barwinska

title: Overcoming catastrophic forgetting in neural networks

venue: Proceedings of the National Academy of Sciences,

authors:
F. Kunstner
L. Balles
P. Hennig

title: Limitations of the empirical fisher approximation for natural gradient descent

venue: arXiv preprint arXiv:1905.12558

authors:
Z. Li
D. Hoiem

title: Learning without forgetting

venue: IEEE Transactions on Pattern Analysis and Machine Intelligence,

authors:
N. Loo
S. Swaroop
R.E. Turner

title: Generalized variational continual learning. arXiv preprint arXiv:2011.12328

authors:
H. Lütcke
D.J. Margolis
F. Helmchen

title: Steady or changing? long-term monitoring of neuronal population activity

venue: Trends in Neurosciences,

authors:
V. Mante
D. Sussillo
K.V. Shenoy
W.T. Newsome

title: Context-dependent computation by recurrent dynamics in prefrontal cortex

title: New insights and perspectives on the natural gradient method

venue: arXiv preprint arXiv:1412.1193

authors:
J. Martens
R. Grosse

title: Optimizing neural networks with Kronecker-factored approximate curvature

venue: In ICML,

authors:
C.V. Nguyen
Y. Li
T.D. Bui
R.E. Turner

title: Variational continual learning

venue: arXiv preprint arXiv:1710.10628

authors:
K. Osawa
S. Swaroop
A. Jain
R. Eschenhagen
R.E. Turner
R. Yokota
M.E. Khan

title: Practical deep learning with Bayesian principles

venue: arXiv preprint arXiv:1906.02506

authors:
P. Pan
S. Swaroop
A. Immer
R. Eschenhagen
R.E. Turner
M.E. Khan

title: Continual deep learning by functional regularisation of memorable past

venue: arXiv preprint arXiv:2004.14070

authors:
A.J. Peters
S.X. Chen
T. Komiyama

title: Emergence of reproducible spatiotemporal activity during motor learning

authors:
Rebuffi
S.-A
A. Kolesnikov
G. Sperl
C.H. Lampert

title: icarl: Incremental classifier and representation learning

authors:
H. Ritter
A. Botev
D. Barber

title: Online structured laplace approximations for overcoming catastrophic forgetting

venue: arXiv preprint arXiv:1805.07810

authors:
D. Rolnick
A. Ahuja
J. Schwarz
T. Lillicrap
G. Wayne

title: Experience replay for continual learning

venue: In Advances in Neural Information Processing Systems,

authors:
M.E. Rule
T. O’Leary
C.D. Harvey

title: Causes and consequences of representational drift

venue: Current Opinion in Neurobiology,

authors:
G. Saha
I. Garg
K. Roy

title: Gradient projection memory for continual learning. arXiv preprint arXiv:2103.09762

authors:
J. Schwarz
W. Czarnecki
J. Luketina
A. Grabska-Barwinska
Y.W. Teh
R. Pascanu
R. Hadsell

title: Progress & compress: A scalable framework for continual learning

venue: In International Conference on Machine Learning,

authors:
H. Shin
J.K. Lee
J. Kim

title: Continual learning with deep generative replay

venue: arXiv preprint arXiv:1705.08690

authors:
M.K. Titsias
J. Schwarz
Matthews
A.G. d. G
R. Pascanu
Y.W. Teh

title: Functional regularisation for continual learning with gaussian processes

venue: In International Conference on Learning Representations

authors:
H. Tseran
M.E. Khan
T. Harada
T.D. Bui

title: Natural variational continual learning

venue: In Continual Learning Workshop NeurIPS,

authors:
G.M. van de Ven
H.T. Siegelmann
A.S. Tolias

title: Brain-inspired replay for continual learning with artificial neural networks

venue: Nature Communications,

authors:
G.M. van de Ven
A.S. Tolias

title: Generative replay with feedback connections as a general strategy for continual learning

venue: arXiv preprint arXiv:1809.10635

title: Three scenarios for continual learning

venue: arXiv preprint arXiv:1904.07734

authors:
J. von Oswald
C. Henning
J. Sacramento
B.F. Grewe

title: Continual learning with hypernetworks

venue: arXiv preprint arXiv:1906.00695

authors:
T. Xu
X. Yu
A.J. Perlik
W.F. Tobin
J.A. Zweig
K. Tennant
T. Jones
Y. Zuo

title: Rapid formation and selective stabilization of synapses for enduring motor

authors:
G. Yang
F. Pan
Gan
W.-B

title: Stably maintained dendritic spines are associated with lifelong memories

authors:
G.R. Yang
M.R. Joglekar
H.F. Song
W.T. Newsome
Wang
X.-J

title: Task representations in neural networks trained to perform many cognitive tasks

authors:
B.M. Yu
J.P. Cunningham
G. Santhanam
S.I. Ryu
K.V. Shenoy
M. Sahani

title: Gaussian-process factor analysis for low-dimensional single-trial analysis of neural population activity

venue: Journal of Neurophysiology,

authors:
G. Zeng
Y. Chen
B. Cui
S. Yu

title: Continual learning of context-dependent processing in neural networks

venue: Nature Machine Intelligence,

authors:
F. Zenke
B. Poole
S. Ganguli

title: Continual learning through synaptic intelligence

heading: 1 Introduction

text:
Catastrophic forgetting is a common feature of machine learning algorithms where training on a new task often leads to poor performance on previously learned tasks. This is in contrast to biological agents which are capable of learning many different behaviors over the course of their lives with little to no interference across tasks. The study of continual learning in biological networks may therefore help inspire novel approaches in machine learning, while the development and study of continual learning algorithms in artificial agents can help us better understand how this challenge is overcome in the biological domain. This is particularly true for more challenging continual learning settings where task identity is not provided at test time, and for continual learning in recurrent neural
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
networks (RNNs), which is important due to the practical and biological relevance of RNNs. However, continual learning in these settings has recently proven challenging for many existing algorithms, particularly those that rely on parameter regularization to mitigate forgetting [12, 13, 47]. In this work, we address these shortcomings by developing a continual learning algorithm that not only encourages good performance across tasks at convergence but also regularizes the optimization path itself using trust region optimization. This leads to improved performance compared to existing methods.
Previous work has addressed the challenge of continual learning in artificial agents using weight regularization, where parameters important for previous tasks are regularized to stay close to their previous values [1, 17, 24, 32, 37, 54]. This approach can be motivated by findings in the neuroscience literature of increased stability for a subset of synapses after learning [49, 50]. More recently, approaches based on projecting gradients into subspaces orthogonal to those that are important for previous tasks have been developed in both feedforward [40, 53] and recurrent [12] neural networks. This is consistent with experimental findings that neural dynamics often occupy orthogonal subspaces across contexts in biological circuits [3, 14, 20, 22]. While these methods have been found to perform well in many continual learning settings, they also suffer from several shortcomings. In particular, while Bayesian weight regularization provides a natural way to weigh previous and current task information, this approach can fail in practice due to its approximate nature and often requires additional tuning of the importance of the prior beyond what would be expected in a rigorous Bayesian treatment [46]. In contrast, while projection-based methods have been found empirically to mitigate catastrophic forgetting, it is unclear how the ‘important subspaces’ should be selected and how such methods behave when task demands begin to saturate the network capacity.
In this work, we develop natural continual learning (NCL), a new method that combines (i) Bayesian continual learning using weight regularization with (ii) an optimization procedure that relies on a trust region constructed from an approximate posterior distribution over the parameters given previous tasks. This encourages parameter updates predominantly in the null-space of previously acquired tasks while maintaining convergence to a maximum of the Bayesian approximate posterior. We show that NCL outperforms previous continual learning algorithms in both feedforward and recurrent networks. We also show that the projection-based methods introduced by Duncker et al. [12] and Zeng et al. [53] can be viewed as approximations to such trust region optimization using the posterior from previous tasks. Finally, we use tools from the neuroscience literature to investigate how the learned networks overcome the challenge of continual learning. Here, we find that the networks learn latent task representations that are stable over time after initial task learning, consistent with results from biological circuits.

heading: 2 Method

text:
Notations We use X>, X−1, Tr(X) and vec(X) to denote the transpose, inverse, trace, and column-wise vectorization of a matrix X . We use X ⊗ Y to represent the Kronecker product between matrices X ∈ Rn×n and Y ∈ Rm×m such that (X ⊗ Y )mi+k,mj+l = XijYkl. We use bold lower-case letters x to denote column vectors. Dk refers to a ‘dataset’ corresponding to task k, which in this work generally consists of a set of input-output pairs {x(i)k ,y (i) k } such that
`k(θ) := log p(Dk|θ) = ∑ i log pθ(y (i) k |x (i) k ) is the task-related performance on task k for a model with parameters θ. Finally, we use D̂k to refer to a dataset generated by inputs from the kth task where {ŷ(i)k ∼ pθ(y|x (i) k )} are drawn from the model distributionM.

heading: 2.1 Bayesian continual learning

text:
Problem statement In continual learning, we train a model on a set of K tasks {D1, . . . ,DK} that arrive sequentially, where the data distribution Dk for task k in general differs from D6=k. The aim is to learn a probabilistic model p(D|θ) that performs well on all tasks. The challenge in the continual learning setting stems from the sequential nature of learning, and in particular from the common assumption that the learner does not have access to “past” tasks (i.e., Dj for j < k) when learning task k. While we enforce this stringent condition in this paper, our approach may be easily combined with memory-based techniques such as coresets or generative replay [8, 10, 13, 32, 34, 36, 38, 42, 43, 45, 48].
Bayesian approach The continual learning problem is naturally formalized in a Bayesian framework whereby the posterior after k−1 tasks is used as a prior for task k. More specifically, we choose a prior p(θ) on the model parameters and compute the posterior after observing k tasks according to Bayes’ rule:
p(θ|D1:k) ∝ p(θ) k∏
k′=1
p(Dk′ |θ)
∝ p(θ|D1:k−1)p(Dk|θ), (1)
where D1:k is a concatenation of the first k tasks (D1, . . . ,Dk). In theory, it is thus possible to compute the exact posterior p(θ|D1:k) after k tasks, while only observing Dk, by using the posterior p(θ|D1:k−1) after k − 1 tasks as a prior. However, as is often the case in Bayesian inference, the difficulty here is that the posterior is typically intractable. To address this challenge, it is common to perform approximate online Bayesian inference. That is, the posterior p(θ|D1:k−1) is approximated by a parametric distribution with parameters φk−1. The approximate posterior q(θ;φk−1) is then used as a prior for task k.
Online Laplace approximation A common approach is to use the Laplace approximation whereby the posterior p(θ|D1:k−1) is approximated as a multivariate Gaussian q using local gradient information [17, 24, 37]. This involves (i) finding a mode µk of the posterior during task k, and (ii) performing a second-order Taylor expansion around µk to construct an approximate Gaussian posterior q(θ;φk) = N (θ;µk,Λ−1k ), where Λk is the precision matrix and φk = (µk,Λk). In this case, gradient-based optimization is used to find the posterior mode on task k (c.f. Equation 1):
µk = arg max θ
log p(θ|Dk,φk−1) (2)
= arg max θ
log p(Dk|θ) + log q(θ;φk−1) (3)
= arg max θ
`k(θ)− 1
2 (θ − µk−1)>Λk−1(θ − µk−1)︸ ︷︷ ︸
:= Lk(θ)
(4)
The precision matrix Λk is given by the Hessian of the negative log posterior at µk: Λk = − ∇2θ log p(θ|Dk,φk−1) ∣∣ θ=µk = H(Dk,µk) + Λk−1, (5)
where H(Dk,µk) = − ∇2θ log p(Dk|θ) ∣∣ θ=µk is the Hessian of the negative log likelihood of Dk.
Continual learning with the online Laplace approximation thus involves two steps for each new task Dk. First, given Dk and the previous posterior q(θ;µk−1,Λ−1k−1) (i.e. the new prior), µk is found using gradient-based optimization (Equation 4). This step can be interpreted as optimizing the likelihood of Dk while penalizing changes in the parameters θ according to their importance for previous tasks, as determined by the prior precision matrix Λk−1. Second, the new posterior precision matrix Λk is computed according to Equation 5.
Approximating the Hessian In practice, computing Λk presents two major difficulties. First, because q(θ;φk) is a Gaussian distribution, Λk has to be positive semi-definite (PSD), which is not guaranteed for the Hessian H(Dk,µk). Second, if the number of model parameters nθ is large, it may be prohibitive to compute a full (nθ × nθ) matrix. To address the first issue, it is common to approximate the Hessian with the Fisher information matrix (FIM; 17, 30, 37):
Fk = Ep(D̂k|θ) [ ∇θ log p(D̂k|θ)∇θ log p(D̂k|θ)> ]∣∣∣ θ=µk ≈ H(Dk,µk) (6)
The FIM is PSD, which ensures that Λk = ∑k k′=1 Fk′ is also PSD. Computing Fk may still be impractical if there are many model parameters, and it is therefore common to further approximate the FIM using structured approximations with fewer parameters. In particular, a diagonal approximation to Fk recovers Elastic Weight Consolidation (EWC; 24), while a Kronecker-factored approximation [31] recovers the method proposed by Ritter et al. [37]. We denote this method ‘KFAC’ and use it in Section 3 as a comparison for our own Kronecker-factored method.

heading: 2.2 Natural continual learning

text:
While the online Laplace approximation has been applied successfully in several continual learning settings [24, 37], it has also been found to perform sub-optimally on a range of problems [12, 46]. Additionally, its Bayesian interpretation in theory prescribes a unique way of weighting the contributions of previous and current tasks to the loss. However, to perform well in practice, weight regularization approaches have been found to require ad-hoc re-weighting of the prior term by several orders of magnitude [24, 37, 46]. These shortcomings could be due to an inadequacy of the approximations used to construct the posterior (Section 2.1). However, we show in Figure 1 that standard gradient descent on the Laplace posterior has important drawbacks even in the exact case. First, we show that exact Bayesian inference on a simple continual regression problem can produce indirect optimization paths along which previous tasks are transiently forgotten as a new task is being learned (Figure 1A; green). Second, when the loss is non-convex, we show that exact Bayesian inference can still lead to catastrophic forgetting (Figure 1B; green).
An alternative approach that has found recent success in a continual learning setting involves projection based methods which restrict parameter updates to a subspace that does not interfere with previous tasks [12, 53]. However, it is not immediately obvious how this projected subspace should be selected in a way that appropriately balances learning on previous and current tasks. Additionally, such projection-based algorithms have fixed points that are minima of the current task, but not necessarily minima of the (negative) Bayesian posterior. This can lead to catastrophic forgetting in the limit of long training times (Figure 1; pink), unless the learning rate is exactly zero in directions that interfere with previous tasks.
To combine the desirable features of both classes of methods, we introduce “Natural Continual Learning” (NCL) – an extension of the online Laplace approximation that also restricts parameter updates to directions which do not interfere strongly with previous tasks. In a Bayesian setting, we can conveniently express what is meant by such directions in terms of the prior precision matrix Λ. In particular, ‘flat’ directions of the prior (low precision) correspond to directions that will not significantly affect the performance on previous tasks. Formally, we derive NCL as the solution of a trust region optimization problem. This involves minimizing the posterior loss Lk(θ) within a region of radius r centered around θ with a distance metric of the form d(θ,θ + δ) = √ δ>Λk−1δ/2 that takes into account the curvature of the prior via its precision matrix Λk−1:
δ = arg min δ
Lk(θ) +∇θLk(θ)>δ subject to 1
2 δ>Λk−1δ ≤ r2, (7)
where Lk(θ + δ) ≈ Lk(θ) + ∇θLk(θ)>δ is a first-order approximation to the updated Laplace objective. The solution to this subproblem is given by δ ∝ Λ−1k−1∇θ`k(θ) − (θ − µk−1) (see
Appendix A for a derivation), which gives rise to the NCL update rule θ ← θ + γ [ Λ−1k−1∇θ`k(θ)− (θ − µk−1) ] (8)
for a learning rate parameter γ (which is implicitly a function of r in Equation 7). To get some intuition for this learning rule, we note that Λ−1k−1 acts as a preconditioner for the first (likelihood) term, which drives learning on the current task while encouraging parameter changes predominantly in directions that do not interfere with previous tasks. Meanwhile, the second term encourages θ to stay close to µk−1, the optimal parameters for the previous task. As we illustrate in Figure 1, this combines the desirable features of both Bayesian weight regularization and projection-based methods. In particular, NCL shares the fixed points of the Bayesian posterior while also mitigating intermediate or complete forgetting of previous tasks by preconditioning with the prior covariance. Notably, if the loss landscape is non-convex (as it generally will be), NCL can converge to a different local optimum from standard weight regularization despite having the same fixed points (Figure 1B).
Implementation The general NCL framework can be applied with different approximations to the Fisher matrix Fk in Equation 6 (see Section 2.1). In this work, we use a Kronecker-factored approximation [31, 37]. However, even after making a Kronecker-factored approximation to Fk for each task k, it remains difficult to compute the inverse of a sum of k Kronecker products (c.f. Equation 5). To address this challenge, we derived an efficient algorithm for making a Kroneckerfactored approximation to Λk = Fk + Λk−1 ≈ Ak ⊗Gk when Λk−1 = Ak−1 ⊗Gk−1 and Fk are also Kronecker products. This approximation minimizes the KL-divergence between N (µk, (Ak ⊗ Gk)
−1) and N (µk, (Λk−1 + Fk)−1) (see Appendix G for details). Before training on the first task, we assume a spherical Gaussian prior θ ∼ N (0, p−2w I). The scale parameter pw can either be set to a fixed value (e.g. 1) or treated as a hyperparameter, and we optimize pw explicitly for our experiments in feedforward networks. NCL also has a parameter α which is used to stabilize the matrix inversion Λ−1k−1 ≈ (Ak−1 ⊗Gk−1 + α2I)−1 (Appendix E). This is equivalent to a hyperparameter used for such matrix inversions in OWM [53] and DOWM [12], and it is important for good performance with these methods. The pw and α are largely redundant for NCL, and we generally prefer to fix α to a small value (10−10) and optimize the pw only. However, for our experiments in RNNs, we instead fix pw = 1 and perform a hyperparameter optimization over α for a more direct comparison with OWM and DOWM. The NCL algorithm is described in pseudocode in Appendix E together with additional implementation and computational details. Finally, while we have derived NCL with a Laplace approximation in this section for simplicity, it can similarly be applied in the variational continual learning framework of Nguyen et al. [32] (Appendix J). Our code is available online1.

heading: 2.3 Related work

text:
As discussed in Section 2.1, our method is derived from prior work that relies on Bayesian inference to perform weight regularization for continual learning [17, 24, 32, 37]. However, we also take inspiration from the literature on natural gradient descent [2, 25] to introduce a preconditioner that encourages parameter updates primarily in flat directions of previously learned tasks (Appendix H).
Recent projection-based methods [12, 40, 53] have addressed the continual learning problem using an update rule of the form θ ← θ + γPL∇θ`k(θ)PR, (9) wherePL andPR are projection matrices constructed from previous tasks which encourage parameter updates that do not interfere with performance on these tasks. Using Kronecker identities, we can rewrite Equation 9 as θ ← θ + γ(PR ⊗ PL)∇θ`k(θ). (10) This resembles the NCL update rule in Equation 8 where we identify PR ⊗PL with the approximate inverse prior precision matrix used for gradient preconditioning in NCL, Λ−1k−1 = A −1 k−1 ⊗G −1 k−1. Indeed, we note that for a Kronecker-structured approximation to Fk, the matrixAk−1 approximates the empirical covariance matrix of the network activations experienced during all tasks up to k − 1 (4, 31, Appendix D), which is exactly the inverse of the projection matrix PR used in previous work [12, 53]. We thus see that NCL takes the form of recent projection-based continual learning algorithms with two notable differences:
1https://github.com//tachukao/ncl
(i) NCL uses a left projection matrix PL designed to approximate the posterior covariance of previous tasks Λ−1k−1 ≈ PR ⊗PL (i.e., the prior covariance on task k; Appendix D), while Zeng et al. [53] use the identity matrix I and Duncker et al. [12] use the covariance of recurrent inputs (Appendix F). Notably, both of these choices of PL still provide reasonable approximations to Λ−1k−1, and thus the parameter updates of OWM and DOWM can also be viewed as projecting out steep directions of the prior on task k (Appendix F). (ii) NCL includes an additional regularization term (θ − µk−1) derived from the Bayesian posterior objective, while Duncker et al. [12] and Zeng et al. [53] do not use such regularization. Importantly, this means that while NCL has a similar preconditioner and optimization path to these projection based methods, NCL has stationary points at the modes of the approximate Bayesian posterior while the stationary points of OWM and DOWM do not incorporate prior information from previous tasks (c.f. Figure 1).
It is also interesting to note that previous Bayesian continual learning algorithms include a hyperparameter λ that scales the prior compared to the likelihood term for the current task [27]:
L(λ)k (θ) = log p(Dk|θ)− λ(θ − µk−1) >Λk−1(θ − µk−1). (11)
To minimize this loss and thus find a mode of the approximate posterior, it is common to employ pseudo-second-order stochastic gradient-based optimization algorithms such as Adam [23] that use their own gradient preconditioner based on an approximation to the Hessian of Equation 11. Interestingly, this Hessian is given by Hk = −H(Dk,θ) − λΛk−1, which in the limit of large λ becomes increasingly similar to preconditioning with the prior precision as in NCL. Consistent with this, previous work using the online Laplace approximation has found that large values of λ are generally required for good performance [24, 37, 46]. Recent work has also combined Bayesian continual learning with natural gradient descent [33, 44], and in this case a relatively high value of λ = 100 was similarly found to maximize performance [33].

heading: 3 Experiments and results

heading: 3.1 NCL in feedforward networks

text:
To verify the utility of NCL for continual learning, we first compared our algorithm to standard methods in feedforward networks across two continual learning benchmarks: split MNIST and split CIFAR-100 (see Appendix B for task details). For each benchmark, we considered three continual learning settings [47]. In the ‘task-incremental’ setting, task identity is available to the network at test time, in our case via a multi-head output layer [5]. In the ‘domain-incremental’ setting, task identity is unavailable at test time, and the output layer is shared between all tasks. Finally, in the ‘class-incremental’ setting, the network has to both infer task identity and solve the task, in our case by performing classification over all possible classes irrespective of which task the input in question is drawn from.
van de Ven and Tolias previously showed that parameter regularization methods such as EWC perform poorly in the domain- and class-incremental settings [47]. We therefore applied NCL as well as synaptic intelligence [SI; 54], online EWC [41], Kronecker factored EWC [KFAC; 37], and orthogonal weight modification [OWM; 53] to split MNIST and split CIFAR-100 in the task-, domain- and class-incremental learning settings. For these continual learning problems, we found that NCL outperformed all the baseline methods in the task- and domain-incremental learning settings (Figure 2). In the class-incremental settings, we found that NCL performed comparably to but slightly worse than OWM. However, both OWM and NCL comfortably outperformed the other compared methods in this setting. These results suggest that the subpar performance of parameter regularization methods can be alleviated by regularizing their optimization paths, particularly in the domain- and class-incremental learning settings.
For the split MNIST and split CIFAR-100 experiments, each baseline method had a single hyperparameter (c for SI, λ for EWC and KFAC, α for OWM, and pw for NCL; Appendix E) that was optimized on a held-out seed (see Appendix I.2). However, by setting the NCL prior to a unit Gaussian, we were also able to achieve good performance across task sets in a hyperparameter-free setting, further highlighting the robustness of the method (see “NCL (no opt)” in Figure 2).

heading: 3.2 NCL in recurrent neural networks

text:
We then proceeded to consider how NCL compares to previous methods in recurrent neural networks (RNNs), a setting that has recently proven challenging for continual learning [12, 13] and which is of interest to the study of continual learning in biological circuits [12, 51]. In these experiments, the task identity is available to the RNN (i.e., we consider the task-incremental learning setting).
Stimulus-response tasks In this section, we consider a set of neuroscience inspired ‘stimulusresponse’ (SR) tasks (51; details in Appendix B). We first compared the performance and behavior of NCL to OWM, the top performing method in the feedforward setting (Figure 2), and to the projection-based DOWM method designed explicitly for RNNs [12]. For a more direct comparison with OWM and DOWM, we fixed the NCL prior to a unit Gaussian for all RNN experiments and instead performed a hyperparameter optimization over ‘α’ used to regularize the matrix inversions for all three methods (Section 2.2, Appendix E, Appendix I.2, 12, 53). Following previous work, we trained RNNs with 256 recurrent units to sequentially solve six stimulus-response tasks [12, 50]. While NCL, OWM and DOWM all managed to learn the six tasks without catastrophic forgetting, we found that NCL achieved superior average performance across tasks after training (Figure 3A).
We then compared NCL, OWM, and DOWM to KFAC, the top performing parameter regularization method in our feedforward experiments (Figure 2) which uses Adam [23] to optimize the objective in Equation 4 with a Kronecker-factored approximation to the posterior precision matrix (Section 2.1; 37). Consistent with the results shown in Duncker et al. [12], we found that NCL, OWM, and DOWM outperformed KFAC with λ = 1 (Figure 3A; see also 12 for a comparison of DOWM and EWC). We note that NCL and KFAC optimize the same objective function (Equation 4) and approximate the posterior precision matrix in the same way, but they differ in the way they precondition the gradient of the objective. These results thus demonstrate empirically that the choice of optimization algorithm is important to prevent forgetting, consistent with the intuition provided by Figure 1.
In feedforward networks, poor performance with weight regularization approaches such as EWC and KFAC has been mitigated by optimizing the hyperparameter λ, which increases the importance of the prior term compared to a standard Bayesian treatment (Equation 11; Section 3.1, 24, 27, 37). We confirmed this here by performing a grid search over λ, which showed that KFAC with λ ∈
[100, 1000] could perform comparably to the projection-based methods (Appendix I.1; Figure 3A). We hypothesize that the good performance provided by high λ is partly due to the approximate second order nature of Adam which, together with the relative increase in the prior term compared to the data term, leads to preconditioning with a matrix resembling the prior Λk−1 (Section 2.3). In support of this hypothesis, we found that the KL divergence between the Adam preconditioner and the approximate prior precision Λk−1 decreased with increasing λ, and that the performance of KFAC with Adam could also be rescued by increasing λ only when computing the preconditioner while retaining λ = 1 when computing the gradients (Appendix I.1).
Stroke MNIST One way to challenge the continual learning algorithms further is to increase the number of tasks. We thus considered an augmented version of the stroke MNIST dataset [SMNIST; 9]. The original dataset consists of the MNIST digits transformed into pen strokes with the direction of the stroke at each time point provided as an input to the network. Similar to Ehret et al. [13], we constructed a continual learning problem by considering consecutive binary classification tasks inspired by the split MNIST task set. We further increased the number of tasks by including a set of extra digits where the x and y dimensions have been swapped in the input stroke data, and another set where both the x and y dimensions have changed sign. We also added high-variance noise to the inputs to increase the task difficulty. This gave rise to a total of 15 binary classification tasks, each with unique digits not used in other tasks, which we sought to learn in a continual fashion using an RNN with 30 recurrent units (see Appendix B for details).
As for the SR task set in Section 3.2, we found that NCL outperformed previous projection-based methods (Figure 3B). We again found that weight regularization with a KFAC approximation performed poorly with λ = 1, and that this poor performance could be partially rescued by optimizing over λ (Figure 3B). To investigate how the difference in performance between NCL and DOWM was affected by their different approximations to the Fisher matrix (Appendix F), we implemented NCL using the DOWM projection matrices as an alternative approximation to the inverse Fisher matrix. We refer to this method as Laplace-DOWM. We then considered how the performance on each task at the end of training depended on task number, averaged over different task permutations (Figure 3C). We found that while Laplace-DOWM outperformed NCL on the first task, this method generally performed worse on subsequent tasks. Notably, Laplace-DOWM exhibited a near-monotonic decrease in relative performance with task number, which is consistent with the intuition that DOWM overestimates the dimensionality of the parameter subspace that matters for previous tasks (Appendix F). In contrast, although neural circuits are known to use orthogonal subspaces in different contexts, there is no general sense that learning more tasks in the past should systematically hinder learning in future contexts for biological agents.

heading: 3.3 Dissecting the dynamics of networks trained on the SMNIST task set

text:
To further investigate how the trained RNNs solve the continual learning problems and how this relates to the neuroscience literature, we dissected the dynamics of networks trained on the SMNIST task set using the NCL algorithm. To do this, we analyzed latent representations of the RNN activity trajectories, as is commonly done to study the collective dynamics of artificial and biological networks [15, 18, 20, 29, 52]. We considered two consecutive classification tasks, namely classifying 4’s vs 5’s (k = 2) and classifying 1’s vs 7’s (k = 3). For each of these tasks, we trained a factor analysis model right after the task was learned, using network activity collected while presenting 50 examples of each of the two input digits associated with the task. We then tracked the network responses to the same set of stimuli at various stages of learning, both before and after the task in question was acquired, using the trained factor analysis model to visualize low-dimensional summaries of the dynamics (Figure 4).
Consistent with the network having successfully learned to solve these two tasks, we found that latent trajectories diverged over time for the two types of inputs in each task. Critically, these diverging dynamics only emerged after the task was learned, and remained highly stable thereafter (Figure 4). The stability of the task-associated representations is consistent with recent work in the neuroscience literature showing that, in a primate reaching task, latent neural trajectories remain stable after learning [15]. Since here we have access to the activity of all neurons throughout the task, we proceeded to quantify the source of this stability at the level of single units. The stability of such single-neuron dynamics after learning has recently been a topic of much interest in biological circuits [7, 28, 39]. In the RNNs, we found that the single-unit representations of a given digit changed during learning of the task involving that digit but stabilized after learning, consistent with work in several distinct biological circuits [6, 11, 16, 19, 21, 35]. Similar results were found using the DOWM algorithm, which was explicitly designed to preserve network dynamics on previously learned tasks [12]. Interestingly, the stable task representations learned by NCL and DOWM differed markedly from a network trained with replay for continual learning, which instead led to task representations that continued to change after initial task acquisition (Appendix I.4). This illustrates how different approaches to continual learning can lead to qualitatively different circuit dynamics, and it suggests the use of continual learning in artificial networks as a model system for biological continual learning.

heading: 4 Discussion

text:
In summary, we have developed a new framework for continual learning based on approximate Bayesian inference combined with trust-region optimization. We showed that this framework encompasses recent projection-based methods and found that it performs better than naive weight regularization. This was particularly evident when task identity was not provided at test time and in recurrent neural networks, settings which have previously been challenging for many continual learning algorithms [12, 13, 47]. Furthermore, we showed that our principled probabilistic approach outperforms previous projection-based methods [12, 53], in particular when the number of tasks and their complexity challenges the network’s capacity. Finally, we analyzed the dynamics of the learned RNNs in a sequential binary classification problem, where we found that the latent dynamics adapt to each new task. We also found that the task-associated dynamics were subsequently conserved during further learning, consistent with experimental reports of stable neural representations [11, 15, 19]. Importantly, our results suggest that preconditioning with the prior covariance can lead to improved performance over existing continual learning algorithms. In future work, it will therefore be interesting to apply this idea to other weight regularization approaches such as EWC with a diagonal approximate posterior [24]. Finally, a separate branch of continual learning utilizes replay-like mechanisms to reduce catastrophic forgetting [8, 26, 34, 42, 43, 46]. While our work has focused on weight regularization, such regularization and replay are not mutually exclusive. Instead, these two approaches have been found to further improve robustness to catastrophic forgetting when combined [32, 45].
Impact and limitations While we have shown that NCL represents an important conceptual and methodological advance for continual learning, it also comes with several limitations. One such limitation arises from the relative difficulty of computing the prior Fisher matrix which is needed for our projection step. Indeed the success of methods such as Adam [23] and EWC [24] is due in part to their ease of implementation which facilitates broad applicability. It will therefore be interesting to investigate how approximations such as a running average of a diagonal approximation to the empirical Fisher matrix as used in Adam could facilitate the development of simple yet powerful variants of NCL.
Furthermore, while NCL mitigates the need to overcount the prior from previous tasks via λ as in KFAC, it does introduce two other (largely redundant) hyperparameters in the form of (i) the scale of the prior before the first task, and (ii) the parameter α used to regularize the inversion of the prior Fisher matrix, similar to OWM and DOWM [12, 53]. While α is an important hyperparameter for OWM and DOWM and we also optimize it in the RNN setting for a more direct comparison (Section 3.2), we find it more natural to set this parameter to a constant small value present only for numerical stability (Appendix E). This leaves the prior scale which we optimize explicitly in the feedforward setting (Section 3.1). However, in future work it would be interesting to consider whether a good prior can be determined in a data free manner to make NCL a hyperparameter-free method. Finally, computing the Fisher matrix used for pre-conditioning requires explicit knowledge of task boundaries. In future work, it will therefore be interesting to develop an algorithm similar to NCL which also works for online learning problems with continually changing task distributions.
Addressing these challenges is important since machine learning algorithms increasingly need to be robust to changing data distributions and dynamic task specifications as they become more prominent in our everyday lives. Much work has therefore gone into the development of methods for continual learning in the machine learning community. However, with the increasing prevalence of practical algorithms for continual learning, it also becomes increasingly important that we understand how and why these algorithms work – insights that can also help us understand when they might fail. In this work, we have therefore attempted to shed light on the relationship between recent methods for continual learning as well as developing a new algorithm with a principled probabilistic interpretation that makes its underlying assumptions more explicit. Taken together, we hope that this work will help improve our understanding of methods for continual learning while also providing an avenue for further research to increase the reliability and robustness of future continual learning algorithms.

heading: Acknowledgements

text:
We are grateful to Siddharth Swaroop, Lea Duncker, Laura Driscoll, Naama Kadmon Harpaz, and Yashar Ahmadian for insightful discussions. We thank Siddharth Swaroop and Robert Pinsler for useful comments on the manuscript.
Funding disclosure
KTJ was funded by a Gates Cambridge scholarship. GMV was supported by the Lifelong Learning Machines (L2M) program of the Defence Advanced Research Projects Agency (DARPA) via contract number HR0011-18-2-0025.

title: Natural continual learning: success is a journey, not (just) a destination


summaries: This paper addresses the problem of constructing injective normalizing flows, which connect some low-dimensional space with the data manifold of interest in the high-dimensional space. Typically injective (or rectangular) flows are composed of two square flows with the “upsampling” padding layer between them. Because the volume change term in the case of injective flows is seemingly intractable, current approaches find workarounds by training these two square flows separately step-by-step. Though this decision has certain drawbacks and may result in suboptimal performance. Current work suggests to employ numerical linear algebra methods to make volume change computation tractable. Experiments demonstrate that taking into account volume сhanging term improves the generation results both for synthetic and real-life data. Interestingly, it also improves out of distribution detection score, allowing the model trained on FashionMNItST to assign lower probability for the MNIST data samples.

summaries: The paper addresses the challenge of learning a normalizing flow, under the assumption that the data probability function is concentrated on a low dimensional manifold. Utilizing the manifold assumption introduces a challenge in computing the volume change term. This paper suggests addressing this computational challenge in two ways: i) If the manifold dimension is low enough, it is reasonable to calculate the map differential exactly based on AD forward mode; ii) alternatively, an unbiased stochastic estimator for the log det can be used.

summaries: This paper is motivated by the desire to circumvent a problematic requirement in Normalizing Flows (NFs) on the function $f: \mathbb{R}^d \rightarrow \mathbb{R}^D$ that produces the pushforward measure --- namely, that $f$ must be strictly invertible. This is difficult for generative models, since the dimension of the latent space $d$ is usually much smaller than that of the target data space $D$, i.e. $d << D$. The invertibility requirement centers around the change-of-variables formula, which is crucial for computing the cost function in NFs, and which requires computing the Jacobian of $f^{-1}$. To enable the use of NFs on latent/data spaces where $d << D$, the authors propose "Rectangular" NFs (RNFs). The name comes from the fact that the Jacobian of $f$ is now rectangular in the $d << D$ case. This involves the use of a differential geometric variant of the change-of-variables formula, and is much more computationally taxing, but provides superior density estimation performance to previous work.

summaries: The paper presents new methods for estimating densities using rectangular normalizing flows. Rectangular normalizing flows, unlike squared normalizing flows, allow considering injective instead of bijective flows. This property is exploited for estimating the density of data residing on a low (d) dimensional manifold in R^D, for D>d. Existing works using rectangular flows to estimate densities supported on low dimensional manifolds assume that the manifold is known, and therefore are restricted to tori, spheres, etc. In contrast, this work does not make this assumption and shows that density estimation for injective flows based on ML is tractable. Two methods are proposed, which are based on incorporating and estimating the “Jacobian-transpose-Jacobian” in the objective. One is based on an exact derivation, which is computationally demanding. The other is based on stochastic gradient estimates. Experimental results on a simulation and a couple of datasets are shown, demonstrating the advantage of the proposed methods compared to an existing method as well as the tradeoff between memory and variance incurred by the two proposed methods.

summaries: The task of interest is to simultaneously learn a manifold (with a fixed predefined dimensionality) and the density on it, from training data. The proposed method, rectangular flows, is developed based on an existing work (Brehmer and Cranmer [6]) that used a two-step training procedure (for manifold matching and density matching, respectively). The main contributions include: (1) revealing that the two-step training can be improved by joint training; (2) a new technique to enable computing the Jacobian-transpose-Jacobian term of the joint training objective with better efficiency. Experiments on simulated data and MNIST/FMNIST are conducted to demonstrate the proposed method.

abstractText: Normalizing flows are invertible neural networks with tractable change-of-volume terms, which allow optimization of their parameters to be efficiently performed via maximum likelihood. However, data of interest are typically assumed to live in some (often unknown) low-dimensional manifold embedded in a high-dimensional ambient space. The result is a modelling mismatch since – by construction – the invertibility requirement implies high-dimensional support of the learned distribution. Injective flows, mappings from lowto high-dimensional spaces, aim to fix this discrepancy by learning distributions on manifolds, but the resulting volume-change term becomes more challenging to evaluate. Current approaches either avoid computing this term entirely using various heuristics, or assume the manifold is known beforehand and therefore are not widely applicable. Instead, we propose two methods to tractably calculate the gradient of this term with respect to the parameters of the model, relying on careful use of automatic differentiation and techniques from numerical linear algebra. Both approaches perform end-to-end nonlinear manifold learning and density estimation for data projected onto this manifold. We study the trade-offs between our proposed methods, empirically verify that we outperform approaches ignoring the volume-change term by more accurately learning manifolds and the corresponding distributions on them, and show promising results on out-of-distribution detection. Our code is available at https://github.com/layer6ai-labs/rectangular-flows.

authors:
A.A. Alemi
I. Fischer
J.V. Dillon
K. Murphy

title: Deep variational information bottleneck

venue: ICLR

authors:
A.A. Alemi
I. Fischer
J.V. Dillon

title: Uncertainty in the variational information bottleneck

venue: arXiv preprint arXiv:1807.00906

authors:
A.G. Baydin
B.A. Pearlmutter
A.A. Radul
J.M. Siskind

title: Automatic differentiation in machine learning: a survey

venue: Journal of machine learning research, 18

authors:
J. Behrmann
W. Grathwohl
R.T. Chen
D. Duvenaud
J.-H. Jacobsen

title: Invertible residual networks

venue: International Conference on Machine Learning, pages 573–582. PMLR

authors:
J. Behrmann
P. Vicol
K.-C. Wang
R. Grosse
J.-H. Jacobsen

title: Understanding and mitigating exploding inverses in invertible neural networks

venue: International Conference on Artificial Intelligence and Statistics, pages 1792–1800. PMLR

authors:
Y. Bengio
A. Courville
P. Vincent

title: Representation learning: A review and new perspectives

venue: IEEE transactions on pattern analysis and machine intelligence, 35(8):1798–1828

title: Constrained optimization and Lagrange multiplier methods

venue: Academic press

authors:
J. Brehmer
K. Cranmer

title: Flows for simultaneous manifold learning and density estimation

venue: Advances in Neural Information Processing Systems, volume 33

authors:
R.T.Q. Chen
J. Behrmann
D.K. Duvenaud
J.-H. Jacobsen

title: Residual flows for invertible generative modeling

venue: Advances in Neural Information Processing Systems, volume 32

authors:
H. Choi
E. Jang
A.A. Alemi

title: Waic

venue: but why? generative ensembles for robust anomaly detection. arXiv preprint arXiv:1810.01392

authors:
R. Cornish
A. Caterini
G. Deligiannidis
A. Doucet

title: Relaxing bijectivity constraints with continuously indexed normalising flows

venue: International Conference on Machine Learning, pages 2133–2143. PMLR

authors:
E. Cunningham
R. Zabounidis
A. Agrawal
I. Fiterau
D. Sheldon

title: Normalizing flows across dimensions

venue: arXiv preprint arXiv:2006.13070

authors:
B. Dai
D. Wipf

title: Diagnosing and enhancing vae models

authors:
L. Dinh
D. Krueger
Y. Bengio

title: Nice: Non-linear independent components estimation

venue: arXiv preprint arXiv:1410.8516

authors:
L. Dinh
J. Sohl-Dickstein
S. Bengio

title: Density estimation using real nvp

authors:
C. Durkan
A. Bekasov
I. Murray
G. Papamakarios

title: Neural spline flows

authors:
J. Gardner
G. Pleiss
K.Q. Weinberger
D. Bindel
A.G. Wilson

title: Gpytorch: Blackbox matrix-matrix gaussian process inference with gpu acceleration

venue: Advances in Neural Information Processing Systems, 31

authors:
M.C. Gemici
D. Rezende
S. Mohamed

title: Normalizing flows on riemannian manifolds

venue: arXiv preprint arXiv:1611.02304

authors:
I.J. Goodfellow
J. Pouget-Abadie
M. Mirza
B. Xu
D. Warde-Farley
S. Ozair
A. Courville
Y. Bengio

title: Generative adversarial nets

venue: Proceedings of the 27th International Conference on Neural Information Processing Systems-Volume 2, pages 2672–2680

authors:
I. Han
D. Malioutov
J. Shin

title: Large-scale log-determinant computation through stochastic chebyshev expansions

venue: International Conference on Machine Learning, pages 908–917. PMLR

authors:
M. Heusel
H. Ramsauer
T. Unterthiner
B. Nessler
S. Hochreiter

title: Gans trained by a two time-scale update rule converge to a local nash equilibrium

venue: Advances in Neural Information Processing Systems, volume 30

authors:
C.-W. Huang
R.T. Chen
C. Tsirigotis
A. Courville

title: Convex potential flows: Universal probability distributions with optimal transport and convex optimization

title: A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines

venue: Communications in Statistics-Simulation and Computation, 18(3):1059– 1076

authors:
S. Ioffe
C. Szegedy

title: Batch normalization: Accelerating deep network training by reducing internal covariate shift

venue: International conference on machine learning, pages 448–456. PMLR

title: Use of different Monte Carlo sampling techniques

venue: Rand Corporation

title: Minima of functions of several variables with inequalities as side constraints

venue: M. Sc. Dissertation. Dept. of Mathematics, Univ. of Chicago

authors:
D.P. Kingma
J. Ba

title: Adam: A method for stochastic optimization

authors:
D.P. Kingma
P. Dhariwal

title: Glow: generative flow with invertible 1⇥ 1 convolutions

venue: Advances in Neural Information Processing Systems, volume 31

authors:
D.P. Kingma
M. Welling

title: Auto-encoding variational bayes

authors:
D.P. Kingma
T. Salimans
R. Jozefowicz
X. Chen
I. Sutskever
M. Welling

title: Improved variational inference with inverse autoregressive flow

authors:
I. Kobyzev
S. Prince
M. Brubaker

title: Normalizing flows: An introduction and review of current methods

venue: IEEE Transactions on Pattern Analysis and Machine Intelligence

authors:
S.G. Krantz
H.R. Parks

title: Geometric integration theory

venue: Springer Science & Business Media

authors:
A. Krizhevsky
G. Hinton

title: Learning multiple layers of features from tiny images

authors:
H.W. Kuhn
A. Tucker. W

title: 1951," nonlinear programming,

venue: In Proceedings of the Second Berkeley Symposium on Mathematical Statistics and Probability(University of California Press, Berkeley, CA),

authors:
A. Kumar
B. Poole
K. Murphy

title: Regularized autoencoders via relaxed injective probability flow

venue: International Conference on Artificial Intelligence and Statistics, pages 4292–4301. PMLR

title: The mnist database of handwritten digits

venue: 1998. URL http://yann. lecun. com/exdb/mnist

authors:
G. Loaiza-Ganem
Y. Gao
J.P. Cunningham

title: Maximum entropy flow networks

authors:
I. Loshchilov
F. Hutter

title: Decoupled weight decay regularization

authors:
C. Lu
J. Chen
C. Li
Q. Wang
J. Zhu

title: Implicit normalizing flows

authors:
E. Mathieu
M. Nickel

title: Riemannian continuous normalizing flows

authors:
E. Nalisnick
A. Matsukawa
Y.W. Teh
D. Gorur

title: and B

venue: Lakshminarayanan. Do deep generative models know what they don’t know? ICLR

authors:
Y. Netzer
T. Wang
A. Coates
A. Bissacco
B. Wu
A.Y. Ng

title: Reading digits in natural images with unsupervised feature learning

authors:
J. Nocedal
S. Wright

title: Numerical optimization

authors:
D. Oktay
N. McGreivy
J. Aduol
A. Beatson
R.P. Adams

title: Randomized automatic differentiation

authors:
G. Papamakarios
T. Pavlakou
I. Murray

title: Masked autoregressive flow for density estimation

authors:
G. Papamakarios
E. Nalisnick
D.J. Rezende
S. Mohamed
B. Lakshminarayanan

title: Normalizing flows for probabilistic modeling and inference

venue: arXiv preprint arXiv:1912.02762

authors:
A. Paszke
S. Gross
F. Massa
A. Lerer
J. Bradbury
G. Chanan
T. Killeen
Z. Lin
N. Gimelshein
L. Antiga
A. Desmaison
A. Kopf
E. Yang
Z. DeVito
M. Raison
A. Tejani
S. Chilamkurthy
B. Steiner
L. Fang
J. Bai
S. Chintala

title: Pytorch: An imperative style, high-performance deep learning library

venue: In Advances in Neural Information Processing Systems 32,

title: Intrinsic statistics on riemannian manifolds: Basic tools for geometric measurements

venue: Journal of Mathematical Imaging and Vision, 25(1):127–154

authors:
K.B. Petersen
M.S. Pedersen

title: The matrix cookbook

venue: URL http://www2.imm

authors:
A. Potapczynski
L. Wu
D. Biderman
G. Pleiss
J.P. Cunningham

title: Bias-free scalable gaussian processes via randomized truncations

venue: International Conference on Machine Learning, to appear

authors:
J. Ren
P.J. Liu
E. Fertig
J. Snoek
R. Poplin
M. Depristo
J. Dillon
B. Lakshminarayanan

title: Likelihood ratios for out-of-distribution detection

authors:
D. Rezende
S. Mohamed

title: Variational inference with normalizing flows

venue: International Conference on Machine Learning, pages 1530–1538. PMLR

authors:
D.J. Rezende
G. Papamakarios
S. Racaniere
M. Albergo
G. Kanwar
P. Shanahan
K. Cranmer

title: Normalizing flows on tori and spheres

venue: International Conference on Machine Learning, pages 8083–8092. PMLR

authors:
B.L. Ross
J.C. Cresswell

title: Tractable density estimation on learned manifolds with conformal embedding flows

venue: arXiv preprint arXiv:2106.05275

title: An introduction to the conjugate gradient method without the agonizing pain

authors:
C. Szegedy
W. Liu
Y. Jia
P. Sermanet
S. Reed
D. Anguelov
D. Erhan
V. Vanhoucke
A. Rabinovich

title: Going deeper with convolutions

venue: Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1–9

authors:
H. Xiao
K. Rasul
R. Vollgraf

title: Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms

venue: arXiv preprint arXiv:1708.07747

heading: 1 Introduction

text:
In recent years, Normalizing Flows (NFs) have become a staple of generative modelling, being widely used for density estimation [14, 15, 45, 28, 16], variational inference [52, 30], maximum entropy modelling [37], and more [46, 31]. In density estimation, we typically have access to a set of points living in some high-dimensional space RD. NFs model the corresponding data-generating distribution as the pushforward of a simple distribution on RD – often a Gaussian – through a smooth bijective mapping. Clever construction of these bijections allows for tractable density evaluation and thus maximum likelihood estimation of the parameters. However, as an immediate consequence of this choice, the learned distribution has support homeomorphic to RD; in particular, the resulting distribution is supported on a set of dimension D. This is not a realistic assumption in practice – especially for density estimation – as it directly contradicts the manifold hypothesis [6] which states that high-dimensional data lives on a lower-dimensional manifold embedded in ambient space.
⇤Authors contributed equally.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
A natural idea to circumvent this misspecification is to consider injective instead of bijective flows, which now push forward a random variable on Rd with d < D to obtain a distribution on some d-dimensional manifold embedded in RD. These mappings admit a change-of-variable formula bearing resemblance to that of bijective flows, but unfortunately the volume-change term becomes computationally prohibitive, which then impacts the tractability of maximum likelihood. While there have been recent efforts towards training flows where the resulting distribution is supported on a low-dimensional manifold [18, 53, 8, 35, 40, 12], these approaches either assume that the manifold is known beforehand or propose various heuristics to avoid the change-of-variable computation. Both of these are undesirable, because, while we should expect most high-dimensional data of interest to exhibit low-dimensional structure, this structure is almost always unknown. On the other hand, we argue that avoiding the volume-change term may result in learning a manifold to which it is difficult to properly assign density, and this approach further results in methods which do not take advantage of density evaluation, undermining the main motivation for using NFs in the first place.
We show that density estimation for injective flows based on maximum likelihood can be made tractable. By carefully leveraging forward- and backward-mode automatic differentiation [3], we propose two methods that allow backpropagating through the volume term arising from the injective change-of-variable formula. The first method involves exact evaluation of this term and its gradient which incurs a higher memory cost; the second uses conjugate gradients [43] and Hutchinson’s trace estimator [23] to obtain unbiased stochastic gradient estimates. Unlike previous work, our methods do not need the data manifold to be specified beforehand, but instead simultaneously estimate this manifold along with the distribution on it end-to-end, thus enabling maximum likelihood training to occur. To the best of our knowledge, ours are the first methods to scale backpropagation through the injective volume-change term to ambient dimensions D close to 3,000. We study the trade-off between memory and variance introduced by our methods and show empirical improvements over injective flow baselines for density estimation. We also show that injective flows obtain state-of-the-art performance for likelihood-based Out-of-Distribution (OoD) detection, assigning higher likelihoods to Fashion-MNIST (FMNIST) [57] than to MNIST [36] with a model trained on the former.

heading: 2 Background

heading: 2.1 Square Normalizing Flows

text:
A normalizing flow [52, 15] is a diffeomorphism f̃✓ : RD ! RD parametrized by ✓, that is, a differentiable bijection with differentiable inverse. Starting with a random variable Z ⇠ pZ for a simple density pZ supported on RD, e.g. a standard Gaussian, the change-of-variable formula states that the random variable X := f̃✓(Z) has density pX on RD given by:
pX(x) = pZ ⇣ f̃ 1✓ (x) ⌘ detJ h f̃✓ i ⇣ f̃ 1✓ (x) ⌘ 1 , (1)
where J[·] is the differentiation operator, so that J[f̃✓](f̃ 1✓ (x)) 2 RD⇥D is the Jacobian of f̃✓ (with respect to the inputs and not ✓) evaluated at f̃ 1✓ (x). We refer to this now standard setup as square flows since the Jacobian is a square matrix. The change-of-variable formula is often written in terms of the Jacobian of f̃ 1✓ , but we use the form of (1) as it is more applicable for the next section. NFs are typically constructed in such a way that not only ensures bijectivity, but also so that the Jacobian determinant in (1) can be efficiently evaluated. When provided with a dataset {xi}ni=1 ⇢ RD, an NF models its generating distribution as the pushforward of pZ through f̃✓, and thus the parameters can be estimated via maximum likelhood as ✓⇤ := argmax✓ Pn i=1 log pX(xi).

heading: 2.2 Rectangular Normalizing Flows

text:
As previously mentioned, square NFs unrealistically result in the learned density pX having Ddimensional support. We follow the injective flow construction of Brehmer and Cranmer [8], where a smooth and injective mapping g : Rd ! RD with d < D is constructed. In this setting, Z 2 Rd is the low-dimensional variable used to model the data as X := g (Z). A well-known result from differential geometry [32] provides an applicable change-of-variable formula:
pX(x) = pZ ⇣ g 1 (x) ⌘ detJ[g ]> ⇣ g 1 (x) ⌘ J[g ] ⇣ g 1 (x) ⌘ 1/2 1(x 2M ), (2)
where M := {g (z) : z 2 Rd}. The Jacobian-transpose-Jacobian determinant now characterizes the change in volume from Z to X . We make several relevant observations: (i) The Jacobian matrix J[g ](g 1 (x)) 2 RD⇥d is no longer a square matrix, and we thus refer to these flows as rectangular.
(ii) Note that g 1 : M ! Rd is only properly defined on M and not RD, and pX is now supported on the d-dimensional manifold M . (iii) We write the indicator 1(x 2M ) explicitly to highlight the fact that this density is not a density with respect to the Lebesgue measure; rather, the dominating measure is a Riemannian measure on the manifold M [48]. (iv) One can clearly verify as a sanity check that when d = D, equation (2) reduces to (1).
Since data points x will almost surely not lie exactly on M , we use a left inverse g† : RD ! Rd in place of g 1 such that g † (g (z)) = z for all z 2 Rd, which exists because g is injective. This is
properly defined on RD, unlike g 1 which only exists over M . Equation (2) then becomes:
pX(x) = pZ ⇣ g† (x) ⌘ detJ[g ]> ⇣ g† (x) ⌘ J[g ] ⇣ g† (x) ⌘ 1/2 . (3)
Note that (3) is equivalent to projecting x onto M as x g (g† (x)), and then evaluating the density from (2) at the projected point.
Now, g is injectively constructed as follows:
g = f̃✓ pad h⌘ and g† = h 1 ⌘ pad † f̃ 1✓ , (4)
where f̃✓ : RD ! RD and h⌘ : Rd ! Rd are both square flows, := (✓, ⌘), and pad : Rd ! RD and pad† : RD ! Rd are defined as pad(z) = (z,0) and pad†(z, z0) = z, where 0, z0 2 RD d. Now, M depends only on ✓ and not ⌘, so we write it as M✓ from now on. Applying (3) yields:
pX(x) = pZ ⇣ g† (x) ⌘ detJ[h⌘] ⇣ g† (x) ⌘ 1 detJ[f✓]> ⇣ f†✓ (x) ⌘ J[f✓] ⇣ f†✓ (x) ⌘ 1/2 , (5)
where f✓ = f̃✓ pad and f†✓ = pad † f̃ 1✓ . We include a derivation of (5) in Appendix A, along with a note on why injective transformations cannot be stacked as naturally as bijective ones.
Evaluating likelihoods is seemingly intractable since constructing flows with a closed-form volumechange term is significantly more challenging than in the square case, even if the relevant matrix is now d⇥ d instead of D ⇥D. Brehmer and Cranmer [8] thus propose a two-step training procedure to promote tractability wherein f✓ and h⌘ are trained separately. After observing that there is no term encouraging x 2M✓, and that x 2M✓ () x = g (g† (x)) () x = f✓(f † ✓ (x)), they decide to simply train f✓ by minimizing the reconstruction error to encourage the observed data to lie on M✓:
✓⇤ = argmin ✓
nX
i=1
xi f✓ ⇣ f†✓ (xi) ⌘ 2
2 . (6)
Note that the above requires computing both f✓ and f†✓ , so that f̃✓ should be chosen as a flow allowing fast evaluation of both f̃✓ and f̃ 1✓ . Architectures such as the Real NVP [15] or follow-up work [28, 16] are thus natural choices for f̃✓, while architectures with an autoregressive component [45, 30] should be avoided. Then, since h⌘ does not appear in the challenging determinant term in (5), h⌘ can be chosen as any normalizing flow, and optimization – for a fixed ✓ – can be tractably achieved by maximum likelihood over the lower-dimensional space:
⌘⇤ = argmax ⌘
nX
i=1
n log pZ ⇣ g† (xi) ⌘ log detJ[h⌘] ⇣ g† (xi) ⌘ o . (7)
In practice, gradient steps in ✓ and ⌘ are alternated. This entire procedure circumvents evaluation of the Jacobian-transpose-Jacobian determinant term in (5), but as we show in section 3, avoiding this term by separately learning the manifold and the density on it comes with its downsides. We then show how to tractably estimate this term in section 4.

heading: 3 Related Work and Motivation

text:
Low-dimensional and topological pathologies The mismatch between the dimension of the modelled support and that of the data-generating distribution has been observed throughout the literature in various ways. Dai and Wipf [13] show, in the context of variational autoencoders [29], that using flexible distributional approximators supported on RD to model data living in a low-dimensional manifold results in pathological behavior where the manifold itself is learned, but not the distribution on it. Cornish et al. [11] demonstrate the drawbacks of normalizing flows for estimating the density of topologically-complex data, and provide a new numerically stable method for learning NFs when the support is not homeomorphic to RD. However, this approach still models the support as being D-dimensional. Behrmann et al. [5] show instabilities associated with NFs – particularly a lack of numerical invertibility, as also explained theoretically by Cornish et al. [11]. This is not too surprising, as attempting to learn a smooth invertible function mapping RD to some low-dimensional manifold is an intrinsically ill-posed problem. This body of work strongly motivates the development of models whose support has matching topology – including dimension – to that of the true data distribution.
Manifold flows A challenge to overcome for obtaining NFs on manifolds is the Jacobian-transposeJacobian determinant computation. Current approaches for NFs on manifolds approach this challenge in one of two ways. The first assumes the manifold is known beforehand [18, 53, 40], severely limiting applicability to low-dimensional data where the true manifold can realistically be known. The second group circumvents the computation of the Jacobian-transpose-Jacobian entirely through various heuristics. Kumar et al. [35] use a potentially loose lower bound of the log-likelihood, and do not explicitly enforce injectivity, resulting in a method for which the change-of-variables almost surely does not hold. Cunningham et al. [12] propose to convolve the manifold distribution with Gaussian noise, which results in the model having high-dimensional support. Finally, Brehmer and Cranmer [8] propose the method we described in subsection 2.2, where manifold learning and density estimation are done separately in order to avoid the log determinant computation. Concurrently to our work, Ross and Cresswell [54] proposed a rectangular flow construction which sacrifices some expressiveness but allows for exact likelihood evaluation.
Why optimize the volume-change term? Learning f✓ and h⌘ separately without the Jacobian of f✓ is concerning: even if f✓ maps to the correct manifold, it might unnecessarily expand and contract volume in such a way that makes correctly learning h⌘ much more difficult than it needs to be. Looking ahead to our experiments, Figure 1 exemplifies this issue: the top-middle panel shows the ground truth density on a 1-dimensional circle in R2, and the top-right panel shows the distribution recovered by the two-step method of Brehmer and Cranmer [8]. We can see that, while the manifold is correctly recovered, the distribution on it is not. The bottom-right panel shows the speed at which f✓⇤ maps R to M✓⇤ : the top of the circle, which should have large densities, also has high speeds. Indeed, there is nothing in the objective discouraging f✓ to learn this behaviour, which implies that the corresponding low-dimensional distribution must be concentrated in a small region and thus making it harder to learn. The bottom-middle panel confirms this explanation: the learned low-dimensional distribution (dark red) does not match what it should (i.e. the distribution of {f†✓⇤(xi)} n i=1, in light red). This failure could have been avoided by learning the manifold in a density-aware fashion by including the Jacobian-transpose-Jacobian determinant in the objective.

heading: 4 Maximum Likelihood for Rectangular Flows: Taming the Gradient

heading: 4.1 Our Optimization Objective

text:
We have argued that including the Jacobian-transpose-Jacobian in the optimization objective is sensible. However, as we previously mentioned, (5) corresponds to the density of the projection of x onto M✓. Thus, simply optimizing the likelihood would not result in learning M✓ in such a way that observed data lies on it, only encouraging projected data points to have high likelihood. We thus maximize the log-likelihood subject to the constraint that the reconstruction error should be smaller than some threshold, i.e. ⇤ = argmax Pn i=1 log pX(xi) subject to Pn i=1 ||xi f✓(f † ✓ (xi))|| 2 2  .
In practice, we use the KKT conditions [26, 34] and maximize the Lagrangian [7] instead:
⇤= argmax
nX
i=1
⇢ log pZ ⇣ g† (xi) ⌘ log detJ[h⌘] ⇣ g† (xi) ⌘ 1
2 log det J>✓ (xi)J✓(xi) (8)
xi f✓ ⇣ f†✓ (xi) ⌘ 2
2
,
where we treat > 0 as a hyperparameter rather than , and denote J[f✓](f†✓ (xi)) as J✓(xi) for simplicity. We have dropped the absolute value since J>✓ (xi)J✓(xi) is always symmetric positive definite, since J✓(xi) has full rank by injectivity of f✓. We now make a technical but relevant observation about our objective: since our likelihoods are Radon-Nikodym derivatives with respect to a Riemannian measure on M✓, different values of ✓ will result in different supports and dominating measures. One should thus be careful to compare likelihoods for models with different values of ✓. However, thanks to the smoothness of the objective over ✓, we should expect likelihoods for values of ✓ which are “close enough” to be comparable for practical purposes. In other words, comparisons remain reasonable locally, and the gradient of the volume-change term should contain relevant information to learn M✓ in such a way that also facilitates learning h⌘ on the pulled-back dataset {f†✓ (xi)} n i=1.

heading: 4.2 Optimizing our Objective: Stochastic Gradients

text:
Note that all the terms in (8) are straightforward to evaluate and backpropagate through except for the third one; in this section we show how to obtain unbiased stochastic estimates of its gradient. In what follows we drop the dependence of the Jacobian on xi from our notation and write J✓, with the understanding that the end computation will be parallelized over a batch of xis. We assume access to an efficient matrix-vector product routine, i.e. computing J>✓ J✓✏ can be quickly achieved for any ✏ 2 Rd. We elaborate on how we obtain these matrix-vector products in the next section. It is a well known fact from matrix calculus [49] that:
@
@✓j log det J>✓ J✓ = tr
✓ (J>✓ J✓) 1 @
@✓j J>✓ J✓
◆ , (9)
where tr denotes the trace operator and ✓j is the j-th element of ✓. Next, we can use Hutchinson’s trace estimator [23], which states that for any matrix M 2 Rd⇥d, tr(M) = E✏[✏>M✏] for any Rd-valued random variable ✏ with zero mean and identity covariance matrix. We can thus obtain an unbiased stochastic estimate of our gradient as:
@
@✓j log det J>✓ J✓ ⇡
1
K
KX
k=1
✏>k (J > ✓ J✓)
1 @
@✓j J>✓ J✓✏k, (10)
where ✏1, . . . , ✏K are typically sampled either from standard Gaussian or Rademacher distributions. Naïve computation of the above estimate remains intractable without explicitly constructing J>✓ J✓. Fortunately, the J>✓ J✓✏ terms can be trivially obtained using the given matrix-vector product routine, avoiding the construction of J>✓ J✓, and then @/@✓jJ > ✓ J✓✏ follows by taking the gradient w.r.t. ✓.
There is however still the issue of computing ✏>(J>✓ J✓) 1 = [(J>✓ J✓) 1✏]>. We use conjugate gradients (CG) [43] in order to achieve this. CG is an iterative method to solve problems of the form Au = ✏ for given A 2 Rd⇥d (in our case A = J>✓ J✓) and ✏ 2 Rd; we include the CG algorithm in Appendix B for completeness. CG has several important properties. First, it is known to recover the solution (assuming exact arithmetic) after at most d steps, which means we can evaluate A 1✏. The solution converges exponentially (in the number of iterations ⌧ ) to the true value [55], so often ⌧ ⌧ d iterations are sufficient for accuracy to many decimal places. In practice, if we can tolerate a certain amount of bias, we can further increase computational speed by stopping iterations early. Second, CG only requires a method to compute matrix-vector products against A, and does not require access to A itself. One such product is performed at each iteration, and CG thus requires at most d matrix-vector products, though again in practice ⌧ ⌧ d products usually suffice. This results in O(⌧d2) solve complexity—less than the O(d3) required by direct inversion methods. We denote A 1✏ computed with conjugate gradients as CG(A; ✏). We can then compute the estimator from (10) as:
@
@✓j log det J>✓ J✓ ⇡
1
K
KX
k=1
CG J>✓ J✓; ✏k > @ @✓j J>✓ J✓✏k. (11)
In practice, we implement this term by noting that CG(J>✓ J✓; ✏) >@/@✓jJ>✓ J✓✏ = @/@✓jstop_gradient(CG(J>✓ J✓; ✏) >)J>✓ J✓✏, thereby taking advantage of the stop_gradient operation from Automatic Differentiation (AD) libraries and allowing us to avoid implementing a custom backward pass. We thus compute the contribution of a point x to the training objective as:
log pZ ⇣ g† (x) ⌘ log detJ[h⌘] ⇣ g† (x) ⌘ x f✓ ⇣ f†✓ (x) ⌘ 2
2 (12)
1
2K
KX
k=1
stop_gradient ⇣ CG
J>✓ J✓; ✏k >⌘ J>✓ J✓✏k
which gives the correct gradient estimate when taking the derivative with respect to .
Linear solvers for Jacobian terms We note that linear solvers like CG have been used before to backpropagate through log determinant computations in the context of Gaussian processes [17], and more recently for square NFs with flexible architectures which do not allow for straightforward Jacobian determinant computations [22, 39]. However, none of these methods require the Jacobiantranspose-Jacobian-vector product routine that we derive in the next section, and to the best of our knowledge, these techniques have not been previously applied for training rectangular NFs. We also point out that recently Oktay et al. [44] proposed a method to efficiently obtain stochastic estimates of J✓✏. While their method cannot be used as a drop-in replacement within our framework as it would result in a biased CG output, we believe this could be an interesting direction for future work. Finally, we note that CG has recently been combined with the Russian roulette estimator [25] to avoid having to always iterate d times while maintaining unbiasedness, again in the context of Gaussian processes [50]. We also leave the exploration of this estimator within our method for future work.

heading: 4.3 AD Considerations: The Exact Method and the Forward-Backward AD Trick

text:
In this section we derive the aforementioned routine for vector products against J>✓ J✓, as well as an exact method that avoids the need for stochastic gradients (for a given x) at the price of increased memory requirements. But first, let us ask: why are these methods needed in the first place? There is work using power series to obtain stochastic estimates of log determinants [20, 9], and one might consider using them in our setting. However, these series require knowledge of the singular values of J>✓ J✓, to which we do not have access (constructing J > ✓ J✓ to obtain its singular values would defeat the purpose of using the power series in the first place), and we would thus not have a guarantee that the series are valid. Additionally, they have to be truncated and thus result in biased estimators, and using Russian roulette estimators to avoid bias [9] can result in infinite variance [11]. Finally, these series compute and backpropagate (w.r.t. ✓) through products of the form ✏>(J>✓ J✓)
m✏ for different values of m, which can easily require more matrix-vector products than our methods. Behrmann et al. [4] address some of the issues with power series approximations as the result of controlling Lipschitz constants, although their estimates remain biased and potentially expensive.
Having motivated our approach, we now use commonly-known properties of AD to derive it; we briefly review these properties in Appendix C, referring the reader to Baydin et al. [3] for more detail. First, we consider the problem of explicitly constructing J✓. This construction can then be used to evaluate J>✓ J✓ and exactly compute its log determinant either for log density evaluation of a trained model, or to backpropagate (with respect to ✓) through both the log determinant computation and the matrix construction, thus avoiding having to use stochastic gradients as in the previous section. We refer to this procedure as the exact method. Naïvely, one might try to explicitly construct J✓ using only backward-mode AD, which would require D vector-Jacobian products (vjps) of the form v>J✓ – one per basis vector v 2 RD (and then stacking the resulting row vectors vertically). A better way to explicitly construct J✓ is with forward-mode AD, which only requires d Jacobian-vector products (jvps) J✓✏, again one per basis vector ✏ 2 Rd (and then stacking the resulting column vectors horizontally). We use a custom implementation of forward-mode AD in the popular PyTorch [47] library2 for the exact method, as well as for the forward-backward AD trick described below.
We now explain how to combine forward- and backward-mode AD to obtain efficient matrix-vector products against J>✓ J✓ in order to obtain the tractable gradient estimates from the previous section.
2PyTorch has a forward-mode AD implementation which relies on the “double backward” trick, which is known to be memory-inefficient. See https://j-towns.github.io/2017/06/12/A-new-trick.html for a description.
Note that v := J✓✏ can be computed with a single jvp call, and then J>✓ J✓✏ = [v >J✓]> can be efficiently computed using only a vjp call. We refer to this way of computing matrix-vector products against J>✓ J✓ as the forward-backward AD trick. We summarize both of our gradient estimators in Appendix D. Note that (12) requires K(⌧ + 1) such matrix-vector products, which is seemingly less efficient as it is potentially greater than the d jvps required by the exact method. However, the stochastic method is much more memory-efficient than its exact counterpart when optimizing over ✓: of the K(⌧ + 1) matrix-vector products needed to evaluate (12), only K require gradients with respect to ✓. Thus only K jvps and K vjps, along with their intermediate steps, must be stored in memory over a training step. In contrast, the exact method requires gradients (w.r.t. ✓) for every one of its d jvp computations, which requires storing these computations along with their intermediate steps in memory.
Our proposed methods thus offer a memory vs. variance trade-off. Increasing K in the stochastic method results in larger memory requirements which imply longer training times, as the batch size must be set to a smaller value. On the other hand, the larger the memory cost, the smaller the variance of the gradient. This still holds true for the exact method, which results in exact gradients, at the cost of increased memory requirements (as long as K ⌧ d; if K is large enough the stochastic method should never be used over the exact one). Table 1 summarizes this trade-off.

heading: 5 Experiments

text: We now compare our methods against the two-step baseline of Brehmer and Cranmer [8], and also study the memory vs. variance trade-off. We use the real NVP [15] architecture for all flows, except we do not use batch normalization [24] as it causes issues with vjp computations. We point out that all comparisons remain fair, and we include a detailed explanation of this phenomenon in Appendix E, along with all experimental details in Appendix G. Throughout, we use the abbreviations RNFs-ML for our maximum likelihood training method, RNFs-TS for the two-step method, and RNFs for rectangular NFs in general. For most runs, we found it useful to anneal the likelihood term(s). That is, at the beginning of training we optimize only the reconstruction term, and then slowly incorporate the other terms. This likelihood annealing procedure helped avoid local optima where the manifold is not recovered (large reconstruction error) but the likelihood of projected data is high.

heading: 5.1 Simulated Data

text:
We consider a simulated dataset where we have access to ground truth, which allows us to empirically verify the deficiencies of RNFs-TS. We use a von Mises distribution, which is supported on the one-dimensional unit circle in R2. Figure 1 shows this distribution, along with its estimates from RNFs-ML (exact) and RNFs-TS. As previously observed, RNFs-TS correctly approximate the manifold, but fail to learn the right distribution on it. In contrast we can see that RNFs-ML, by virtue of including the Jacobian-transpose-Jacobian term in the optimization, manage to recover both the manifold and the distribution on it (top left panel), while also resulting in an easier-to-learn lowdimensional distribution (bottom middle panel) thanks to f✓⇤ mapping to M✓⇤ at a more consistent speed (bottom left panel). We do point out that, while the results presented here are representative of usual runs for both methods, we did have runs with different results which we include in Appendix G for completeness. We finish with the observation that even though the line and the circle are not homeomorphic and thus RNFs are not perfectly able to recover the support, they manage to adequately approximate it.
Table 2: FID-like metric for tabular data (lower is better). Bolded runs are the best or overlap with it.
Method POWER GAS HEPMASS MINIBOONE
RNFs-ML (exact) 0.067± 0.016 0.138± 0.023 0.486± 0.032 0.978± 0.082 RNFs-ML (K = 1) 0.083± 0.015 0.110± 0.021 0.779± 0.191 1.001± 0.051 RNFs-ML (K = 10) 0.113± 0.037 0.140± 0.013 0.495± 0.055 0.878± 0.083 RNFs-TS 0.178± 0.024 0.161± 0.016 0.649± 0.081 1.085± 0.062

heading: 5.2 Tabular Data

text: We now turn our attention to the tabular datasets used by Papamakarios et al. [45], now a common benchmark for NFs as well. As previously mentioned, one should be careful when comparing models with different supports, as we cannot rely on test likelihoods as a metric. We take inspiration from the FID score [21], which is commonly used to evaluate quality of generated images when likelihoods are not available. The FID score compares the first and second moments of a well-chosen statistic – taken in practice to be the values of the last hidden layer of a pre-trained inception network [56] – from the model and data distributions using the squared Wasserstein-2 metric (between Gaussians). Here, we take the statistic to be the data itself instead of the final hidden units of a pre-trained classifier: in other words, our metric compares the mean and covariance of generated data against those of observed data with the same squared Wasserstein-2 metric. We include the mathematical formulas for computing both FID and our modified version for tabular data in Appendix F. We use early stopping with our FID-like score across all models. Our results are summarized in Table 2, where we can see that RNFs-ML consistently do a better job at recovering the underlying distribution. Once again, these results emphasize the benefits of including the Jacobian-transpose-Jacobian in the objective. Interestingly, except for HEPMASS, the results from our stochastic version with K = 1 are not significantly exceeded by the exact version or using a larger value of K, suggesting that the added variance does not result in decreased empirical performance. We highlight that no tuning was done (except on GAS for which we changed d from 4 to 2), RNFs-ML outperformed RNFs-TS out-of-the-box here (details are in Appendix G). We report training times in Appendix G, and observe that RNFs-ML take a similar amount of time as RNFs-TS to train for datasets with lower values of D, and while we do take longer to train for the other datasets, our training times remain reasonable and we often require fewer epochs to converge.

heading: 5.3 Image Data and Out-of-Distribution Detection

text:
We also compare RNFs-ML to RNFs-TS for image modelling on MNIST and FMNIST. We point out that these datasets have ambient dimension D = 784, and being able to fit RNFs-ML is in itself noteworthy: to the best of our knowledge no previous method has scaled optimizing the Jacobiantranspose-Jacobian term to these dimensions. We use FID scores both for comparing models and for early stopping during training. We also used likelihood annealing, with all experimental details again given in Appendix G. We report FID scores in Table 3, where we can see that we outperform RNFs-TS. Our RNFs-ML (K = 1) variant also outperforms its decreased-variance counterparts. This initially puzzling behaviour is partially explained by the fact that we used the K = 1 variant to tune the model (being the cheapest one to train), and then used the tuned hyperparameters for a single run of the other two variants. Nonetheless, once again these results suggest that the variance induced by our stochastic method is not empirically harmful, and that while using the exact method should be the default whenever feasible, using K = 1 otherwise is sensible. We also report training times where we can see the computational benefits of our stochastic method, as well as visualizations of samples, in Appendix G.
We also compare performance on the CIFAR-10 dataset [33], for which D = 3,072. Once again, being able to fit RNFs-ML in this setting is in itself remarkable. We do not include RNFs-ML (K = 4) results because of limited experimentation on CIFAR-10 due to computational cost (experimental details, including hyperparameters which we tried, are given in Appendix G). We can see that, while RNFs-TS outperformed RNFs-ML (K = 1) – which we hypothesize might be reversed given more tuning – our RNFs-ML (exact) version is the best performing model, yet again highlighting the importance of including the change-of-volume term in the objective.
We further evaluate the performance of RNFs for OoD detection. Nalisnick et al. [41] pointed out that square NFs trained on FMNIST assign higher likelihoods to MNIST than they do to FMNIST. While there has been research attempting to fix this puzzling behaviour [1, 2, 10, 51], to the best of our knowledge no method has managed to correct it using only likelihoods of trained models. Figure 2 shows that RNFs remedy this phenomenon, and that models trained on FMNIST assign higher test likelihoods to FMNIST than to MNIST. This correction does not come at the cost of strange behaviour now emerging in the opposite direction (i.e. when training on MNIST, see Appendix G for a histogram). Table 3 quantifies these results (arrows point from in-distribution datasets to OoD ones) with the accuracy of a decision stump using only log-likelihood, and we can see that the best-performing RNFs models essentially solve this OoD task. While we leave a formal explanation of this result for future work, we believe this discovery highlights the importance of
properly specifying models and of ensuring the use of appropriate inductive biases, in this case low intrinsic dimensionality of the observed data. The strong performance of RNFs-TS here seems to indicate that this is a property of RNFs rather than of our ML training method specifically, although our exact approach is still used to compute these log-likelihoods at test time. We include additional results on OoD detection using reconstruction errors – along with a discussion – in Appendix G, where we found the opposite unexpected behaviour: FMNIST always has smaller reconstruction errors, regardless of which dataset was used for training.

heading: 6 Scope and Limitations

text:
In this paper we address the dimensionality-based misspecification of square NFs while properly using maximum likelihood as the training objective, thus providing an advancement in the training of RNFs. Our methods however remain topologically misspecified: even though we can better address dimensionality, we can currently only learn manifolds homeomorphic to Rd. For example, one could conceive of the MNIST manifold as consisting of 10 connected components (one per digit), which cannot be learned by f✓. It is nonetheless worth noting that this limitation is shared by other deep generative modelling approaches, for example GANs [19] result in connected supports (since the image of a connected set under a continuous function is connected). We observed during training in image data that the residuals of CG were not close to 0 numerically, even after d steps, indicating poor conditioning and thus possible numerical non-invertibility of the matrix J>✓ J✓. We hypothesize that this phenomenon is caused by topological mismatch, which we also conjecture affects us more than the baseline as our CG-obtained (or from the exact method) gradients might point in an inaccurate direction. We thus expect our methods in particular to benefit from improved research on making flows match the target topology, for example via continuous indexing [11].
Additionally, while we have successfully scaled likelihood-based training of RNFs far beyond current capabilities, our methods – even the stochastic one – remain computationally expensive for higher dimensions, and further computational gains remain an open problem. We also attempted OoD detection on CIFAR-10 against the SVHN dataset [42], and found that neither RNFs-ML nor RNFsTS has good performance, although anecdotally we may have at least improved on the situation outlined by Nalisnick et al. [41]. We hypothesize these results might be either caused by topological mismatch, or corrected given more tuning.

heading: 7 Conclusions and Broader Impact

text: In this paper we argue for the importance of likelihood-based training of rectangular flows, and introduce two methods allowing to do so. We study the benefits of our methods, and empirically show that they are preferable to current alternatives. Given the methodological nature of our contributions, we do not foresee our work having any negative ethical implications or societal consequences.

heading: Acknowledgements

text: We thank Brendan Ross, Jesse Cresswell, and Maksims Volkovs for useful comments and feedback. We would also like to thank Rob Cornish for the excellent CIFs codebase upon which our code is built, and Emile Mathieu for plotting suggestions. GP and JPC are supported by the Simons Foundation, McKnight Foundation, the Grossman Center, and the Gatsby Charitable Trust.

title: Rectangular Flows for Manifold Learning