S PECTRAL I NFERENCE N ETWORKS : U NIFYING D EEP AND S PECTRAL L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
We present Spectral Inference Networks, a framework for learning eigenfunctions of linear operators by stochastic optimization.
Spectral Inference Networks generalize Slow Feature Analysis to generic symmetric operators and are closely related to Variational Monte Carlo methods from computational physics.
As such, they can be a powerful tool for unsupervised representation learning from video or graph-structured data.
We cast training Spectral Inference Networks as a bilevel optimization problem which allows for online learning of multiple eigenfunctions.
We show results of training Spectral Inference Networks on problems in quantum mechanics and feature learning for videos on synthetic datasets.
Our results demonstrate that Spectral Inference Networks accurately recover eigenfunctions of linear operators and can discover interpretable representations from video in a fully unsupervised manner
I NTRODUCTION
Spectral algorithms are central to machine learning and scientiﬁc computing.
In machine learning, eigendecomposition and singular value decomposition are foundational tools, used for PCA as well as a wide variety of other models.
In scientiﬁc applications, solving for the eigenfunction of a given linear operator is central to the study of PDEs, and gives the time-independent behavior of classical and quantum systems.
For systems where the linear operator of interest can be represented as a reasonably-sized matrix, full eigendecomposition can be achieved in O ( n 3 ) time (Pan et al., 1998), and in cases where the matrix is too large to diagonalize completely (or even store in memory), iterative algorithms based on Krylov subspace methods can efﬁciently compute a ﬁxed number of eigenvectors by repeated application of matrix-vector products (Golub & Van Loan, 2012).
At a larger scale, the eigenvectors themselves cannot be represented explicitly in memory.
This is the case in many applications in quantum physics and machine learning, where the state space of interest may be combinatorially large or even continuous and high dimensional.
Typically, the eigenfunctions of interest are approximated from a ﬁxed number of points small enough to be stored in memory, and then the value of the eigenfunction at other points is approximated by use of the
Nystr¨om method (Bengio et al., 2004).
As this depends on evaluating a kernel between a new point and every point in the training set, this is not practical for large datasets, and some form of function approximation is necessary.
By choosing a function approximator known to work well in a certain domain, such as convolutional neural networks for vision, we may be able to bias the learned representation towards reasonable solutions in a way that is difﬁcult to encode by choice of kernel.
In this paper, we propose a way to approximate eigenfunctions of linear operators on highdimensional function spaces with neural networks, which we call Spectral Inference Networks (SpIN) and train these networks via bilevel stochastic optimization.
We show our method ﬁnds correct eigenfunctions of problems in quantum physics and discovers interpretable representations from video.
This signiﬁcantly extends prior work on unsupervised learning without a generative model and we expect will be useful in scaling many applications of spectral methods.
R ELATED W ORK
Spectral methods are mathematically ubiquitous, arising in a number of diverse settings.
Spectral clustering (Ng et al., 2002), normalized cuts (Shi & Malik, 2000) and Laplacian eigenmaps (Belkin
Niyogi, 2002) are all machine learning applications of spectral decompositions applied to graph
Laplacians.
Related manifold learning algorithms like LLE (Tenenbaum et al., 2000) and IsoMap (Roweis & Saul, 2000) also rely on eigendecomposition, with a different kernel.
Spectral algorithms can also be used for asymptotically exact estimation of parametric models like hidden Markov models and latent Dirichlet allocation by computing the SVD of moment statistics (Hsu et al., 2012;
Anandkumar et al., 2012).
In the context of reinforcement learning, spectral decomposition of predictive state representations has been proposed as a method for learning a coordinate system of environments for planning and control (Boots et al., 2011), and when the transition function is symmetric its eigenfunctions are also known as proto-value functions (PVFs) (Mahadevan & Maggioni, 2007).
PVFs have also been proposed by neuroscientists as a model for the emergence of grid cells in the entorhinal cortex (Stachenfeld et al., 2017).
The use of PVFs for discovering subgoals in reinforcement learning has been investigated in (Machado et al., 2017) and combined with function approximation in (Machado et al., 2018), though using a less rigorous approach to eigenfunction approximation than SpIN.
A qualitative comparison of the two approaches is given in the supplementary material in Sec.
C.3.
Spectral learning with stochastic approximation has a long history as well.
Probably the earliest work on stochastic PCA is that of “Oja’s rule” (Oja, 1982), which is a Hebbian learning rule that converges to the ﬁrst principal component, and a wide variety of online SVD algorithms have appeared since.
Most of these stochastic spectral algorithms are concerned with learning ﬁxed-size eigenvectors from online data, while we are concerned with cases where the eigenfunctions are over a space too large to be represented efﬁciently with a ﬁxed-size vector.
The closest related work in machine learning on ﬁnding eigenfunctions by stochastic approximation is Slow Feature Analysis (SFA) (Wiskott & Sejnowski, 2002), which is a special case of SpIN.
SFA is equivalent to function approximation for Laplacian eigenmaps (Sprekeler, 2011), and it has been shown that optimizing for the slowness of features in navigation can also lead to the emergence of grid cells (Wyss et al., 2006; Franzius et al., 2007).
SFA has primarily been applied to train shallow or linear models, and when trained on deep models is typically trained in a layer-wise fashion, rather than end-to-end (Kompella et al., 2012; Sun et al., 2014).
The features in SFA are learned sequentially, from slowest to fastest, while SpIN allows for simultaneous learning of all eigenfunctions, which is more useful in an online setting.
Spectral methods and deep learning have been combined in other ways.
The spectral networks of
Bruna et al.
(2014) are a generalization of convolutional neural networks to graph and manifold structured data based on the idea that the convolution operator is diagonal in a basis deﬁned by eigenvectors of the Laplacian.
In (Ionescu et al., 2015) spectral decompositions were incorporated as differentiable layers in deep network architectures.
Spectral decompositions have been used in combination with the kernelized Stein gradient estimator to better learn implicit generative models like GANs (Shi et al., 2018).
While these use spectral methods to design or train neural networks, our work uses neural networks to solve large-scale spectral decompositions.
In computational physics, the ﬁeld of approximating eigenfunctions of a Hamiltonian operator is known as Variational Quantum Monte Carlo (VMC) (Foulkes et al., 2001).
VMC methods are usually applied to ﬁnding the ground state (lowest eigenvalue) of electronic systems, but extensions to excited states (higher eigenvalues) have been proposed (Blunt et al., 2015).
Typically the class of function approximator is tailored to the system, but neural networks have been used for ground state calculations (Carleo & Troyer, 2017).
Stochastic optimization for VMC dates back at least to
Harju et al.
Most of these methods use importance sampling from a well-chosen distribution to eliminate the bias due to ﬁnite batch sizes.
In machine learning we are not free to choose the distribution from which the data is sampled, and thus cannot take advantage of these techniques.
M ETHOD 3. S PECTRAL D ECOMPOSITION AS O PTIMIZATION
Eigenvectors of a matrix A are deﬁned as those vectors u such that
Au = λ u for some scalar λ , the eigenvalue.
It is also possible to deﬁne eigenvectors as the solution to an optimization problem.
If
A is a symmetric matrix, then the largest eigenvector of A is the solution of:
[Equation]
or equivalently (up to a scaling factor in u )
[Equation]
This is the Rayleigh quotient , and it can be seen by setting derivatives equal to zero that this is equivalent to ﬁnding u such that A u = λ u , where λ is equal to the value of the Rayleigh quotient.
We can equivalently ﬁnd the lowest eigenvector of A by minimizing the Rayleigh quotient instead.
Amazingly, despite being a nonconvex problem, algorithms such as power iteration converge to the global solution of this problem (Daskalakis et al., 2018, Sec. 4).
To compute the top N eigenvectors U = ( u 1 , . .
. , u N ) , we can solve a sequence of maximization problems:
[Equation]
If we only care about ﬁnding a subspace that spans the top N eigenvectors, we can divide out the requirement that the eigenvectors are orthogonal to one another, and reframe the problem as a single optimization problem (Edelman et al., 1998, Sec. 4.4):
[Equation]
or, if u i denotes row i of U :
[Equation]
Note that this objective is invariant to right-multiplication of U by an arbitrary matrix, and thus we do not expect the columns of U to be the separate eigenvectors.
We will discuss how to break this symmetry in Sec. 3.3.
3. F ROM E IGENVECTORS TO E IGENFUNCTIONS
We are interested in the case where both A and u are too large to represent in memory.
Suppose that instead of a matrix A we have a symmetric (not necessarily positive deﬁnite) kernel k ( x , x (cid:48) ) where x and x (cid:48) are in some measurable space Ω , which could be either continuous or discrete.
Let the inner product on Ω be deﬁned with respect to a probability distribution with density p ( x ) , so that (cid:104) f, g (cid:105) =
(cid:82) f
( x ) g ( x ) p ( x ) d x = E x ∼ p ( x ) [ f ( x ) g ( x )] .
In theory this could be an improper density, such as the uniform distribution over R n , but to evaluate it numerically there must be some proper distribution over Ω from which the data are sampled.
We can construct a symmetric linear operator K from k as K [ f ]( x ) =
E x (cid:48)
[ k ( x , x (cid:48) )
f ( x (cid:48) )] .
To compute a function that spans the top N eigenfunctions of this linear operator, we need to solve the equivalent of Eq. 5 for function spaces.
Replacing rows i and j with points x and x (cid:48) and sums with expectations, this becomes:
[Equation]
where the optimization is over all functions u : Ω → R N such that each element of u is an integrable function under the metric above.
Also note that as u i is a row vector while u ( x ) is a column vector, the transposes are switched.
This is equivalent to solving the constrained optimization problem
[Equation]
For clarity, we will use Σ = E x (cid:2) u ( x ) u ( x ) T (cid:3) to denote the covariance 1 of features and Π =
E x , x (cid:48) (cid:2) k
( x , x (cid:48) ) u ( x ) u ( x (cid:48) )
T (cid:3) to denote the kernel-weighted covariance throughout the paper, so the objective in Eq.
6 becomes Tr( Σ − 1 Π ) .
The empirical estimate of these quantities will be denoted as ˆΣ and ˆ Π .
The form of the kernel k often allows for simpliﬁcation to Eq.
If Ω is a graph, and k ( x , x (cid:48) )
=
−if x (cid:54) = x (cid:48) and are neighbors and 0 otherwise, and k ( x , x ) is equal to the total number of neighbors of x , this is the graph Laplacian , and can equivalently be written as:
[Equation]
for neighboring points (Sprekeler, 2011, Sec. 4.1).
It’s clear that this kernel penalizes the difference between neighbors, and in the case where the neighbors are adjacent video frames this is Slow
Feature Analysis (SFA) (Wiskott & Sejnowski, 2002).
Thus SFA is a special case of SpIN, and the algorithm for learning in SpIN here allows for end-to-end online learning of SFA with arbitrary function approximators.
The equivalent kernel to the graph Laplacian for Ω = R n is
[Equation]
where e i is the unit vector along the axis i .
This converges to the differential Laplacian, and the linear operator induced by this kernel is ∇ 2 (cid:44) (cid:80) i ∂ 2 ∂x 2 i , which appears frequently in physics applications.
The generalization to generic manifolds is the Laplace-Beltrami operator.
Since these are purely local operators, we can replace the double expectation over x and x (cid:48) with a single expectation.
There are many possible ways of solving the optimization problems in Equations 6 and 7.
In principle, we could use a constrained optimization approach such as the augmented Lagrangian method (Bertsekas, 2014), which has been successfully combined with deep learning for approximating maximum entropy distributions (Loaiza-Ganem et al., 2017).
In our experience, such an approach was difﬁcult to stabilize.
We could also construct an orthonormal function basis and then learn some ﬂow that preserves orthonormality.
This approach has been suggested for quantum mechanics problems by Cranmer et al.
But, if the distribution p ( x ) is unknown, then the inner product (cid:104) f, g (cid:105) is not known, and constructing an explicitly orthonormal function basis is not possible. Also, ﬂows can only be deﬁned on continuous spaces, and we are interested in methods that work for large discrete spaces as well.
Instead, we take the approach of directly optimizing the quotient in Eq. 6.
3. L EARNING O RDERED E IGENFUNCTIONS
Since Eq. 6 is invariant to linear transformation of the features u ( x ) , optimizing it will only give a function that spans the top N eigenfunctions of K .
If we were to instead sequentially optimize the
Rayleigh quotient for each function u i ( x ) :
[Equation]
we would recover the eigenfunctions in order.
However, this would be cumbersome in an online setting.
It turns out that by masking the ﬂow of information from the gradient of Eq.
6 correctly, we can simultaneously learn all eigenfunctions in order.
First, we can use the invariance of trace to cyclic permutation to rewrite the objective in Eq.
6 as
Tr
[Equation]
where L is the Cholesky decomposition of Σ .
Let Λ =
L − 1 ΠL − T , this matrix has the convenient property that the upper left n × n block only depends on the ﬁrst n functions u 1: n ( x ) =
( u 1 ( x ) , . . .
, u n ( x )) T .
This means the maximum of (cid:80) ni =1 Λ ii with respect to u 1: n ( x ) spans the ﬁrst n < N eigenfunctions.
If we additionally mask the gradients of Λ ii so they are also independent of any u j ( x ) where j is less than i :
[Equation]
and combine the gradients for each i into a single masked gradient ˜ ∇
u Tr( Λ ) = (cid:80) i ˜
∇ u Λ ii =
( ∂ Λ 11 ∂u 1 , . . .
, ∂ Λ NN ∂u N ) which we use for gradient ascent, then this is equivalent to independently optimizing each u i ( x ) towards the objective Λ ii .
Note that there is still nothing forcing all u ( x ) to be orthogonal.
If we explicitly orthogonalize u ( x ) by multiplication by L − 1 , then we claim that the resulting v ( x ) =
L − 1 u ( x ) will be the true ordered eigenfunctions of K .
A longer discussion justifying this is given in the supplementary material in Sec.
A.
The closed form expression for the masked gradient, also derived in the supplementary material, is given by:
[Equation]
where triu and diag give the upper triangular and diagonal of a matrix, respectively.
This gradient can then be passed as the error from u back to parameters θ , yielding:
[Equation]
To simplify notation we can express the above as
[Equation]
Where J Π ( A ) =
k ( x , x (cid:48) )
u ( x ) T A
∂
u ∂θ and J Σ ( A )
=
u ( x ) T A ∂
u ∂θ are linear operators that denote left-multiplication of the Jacobian of Π and Σ with respect to θ by A .
A TensorFlow implementation of this gradient is given in the supplementary material in Sec.
B.
3. B ILEVEL O PTIMIZATION
The expression in Eq.
14 is a nonlinear function of multiple expectations, so naively replacing Π , Σ , L , Λ and their gradients with empirical estimates will be biased.
This makes learning Spectral
Inference Networks more difﬁcult than standard problems in machine learning for which unbiased gradient estimates are available.
We can however reframe this as a bilevel optimization problem, for which convergent algorithms exist.
Bilevel stochastic optimization is the problem of simultaneously solving two coupled minimization problems min x f ( x , y ) and min y g
( x , y ) for which we only have noisy unbiased estimates of the gradient of each: E
[ F ( x , y )]
=
∇ x f ( x , y ) and E [ G ( x , y )]
= ∇ y g
( x , y ) .
Bilevel stochastic problems are common in machine learning and include actor-critic methods, generative adversarial networks and imitation learning (Pfau & Vinyals, 2016).
It has been shown that by optimizing the coupled functions on two timescales then the optimization will converge to simultaneous local minima of f with respect to x and g with respect to y (Borkar, 1997):
[Equation]
[Equation]
(16) where lim t →∞ α t β t = 0 , (cid:80) t α t = (cid:80)
[Equation]
β
[Equation]
= ∞ , (cid:80) t
α 2 t < ∞ , (cid:80) t β 2 t < ∞ .
By replacing Σ and J Σ with a moving average in Eq.
14, we can cast learning Spectral Inference
Networks as exactly this kind of bilevel problem.
Throughout the remainder of the paper, let ˆ X t denote the empirical estimate of a random variable X from the minibatch at time t , and let ¯ X t represent the estimate of X from a moving average, so ¯ Σ t and ¯ J Σ t are deﬁned as:
[Equation]
[Equation] [Equation]
This moving average is equivalent to solving
[Equation]
1: given symmetric kernel k , decay rates β t , ﬁrst order optimizer O PTIM 2: initialize parameters θ 0 , average covariance ¯ Σ
= I , average Jacobian of covariance ¯ J Σ 0 =3: while not converged do 4:
Get minibatches x t 1 , . .
, x tN and x (cid:48) t 1 , . . .
, x (cid:48) tN 5: ˆ Σ t = 12 (cid:0) 1 N (cid:80) i u θ t ( x ti ) u
θ t ( x ti ) T + 1 N (cid:80) i u θ t ( x (cid:48) ti )
u θ t ( x (cid:48) ti ) T (cid:1) , covariance of minibatches 6: ˆ Π t = 1 N (cid:80) i k
( x ti , x (cid:48) ti )
u θ t ( x ti ) u θ t ( x (cid:48) ti ) T 7: ¯ Σ t ← (1 − β t ) ¯ Σ t − 1 + β t ˆ Σ t 8: ¯ J Σ t ← (1 − β t )¯ J Σ t − 1 + β t
ˆ J Σ t 9: ¯ L t ← Cholesky decomposition of ¯ Σ t 10:
Compute gradient ˜ ∇ θ
Tr( Λ ( ˆ Π t , ¯ Σ t , ˆ J Π t , ¯ J Σ t ))
according to Eq.11: θ t ← O PTIM ( θ t − 1 , ˜ ∇ θ
Tr( Λ ( ˆ Π t , ¯ Σ t , ˆ J Π t , ¯ J Σ t )) 12: result Eigenfunctions v θ ∗ ( x ) =
L − 1 u θ ∗ ( x ) of K [ f ]( x ) =
E x (cid:48)
[ k ( x , x (cid:48) )
f ( x (cid:48) )] by stochastic gradient descent and clearly has the true Σ and J Σ as a minimum for a ﬁxed θ .
Note that Eq.
14 is a linear function of Π and J Π , so plugging in ˆ Π t and ˆ J Π t gives an unbiased noisy estimate.
By also replacing terms that depend on Σ and J Σ with ¯ Σ t and ¯ J Σ t , then alternately updating the moving averages and θ t , we convert the problem into a two-timescale update.
Here θ t corresponds to x t , ¯ Σ t and ¯ J Σ t correspond to y t , ˜ ∇ θ
Tr( Λ ( ˆ Π t , ¯ Σ t , ˆ J Π t , ¯ J Σ t )) corresponds to
F ( x t , y t ) and ( ¯ Σ t − 1 − ˆΣ t , ¯ J Σ t − 1 − ˆ J Σ t ) corresponds to G ( x t , y t ) .
The full algorithm for training Spectral Inference Networks is given in Alg. 1, with TensorFlow pseudocode in the supplementary material in Sec.
B. There are two things to note about this algorithm.
First, we have to compute an explicit estimate ˆ J Σ t of the Jacobian of the covariance with respect to the parameters at each iteration.
That means if we have N eigenfunctions we are computing, each step of training will require N 2 backward gradient computations.
This will be a bottleneck in scaling the algorithm, but we found this approach to be more stable and robust than others.
Secondly, while the theory of stochastic optimization depends on proper learning rate schedules, in practice these proper learning rate schedules are rarely used in deep learning.
Asymptotic convergence is usually less important than simply getting into the neighborhood of a local minimum, and even for bilevel problems, a careful choice of constant learning rates often sufﬁces for good performance.
We follow this practice in our experiments and pick constant values of α and β .
E XPERIMENTS
In this section we present empirical results on a quantum mechanics problem with a known closedform solution, and an example of unsupervised feature learning from video without a generative model.
We also provide experiments comparing our approach against the successor feature approach of Machado et al.
(2018) for eigenpurose discovery on the Arcade Learning Environment in Sec.
C. in the supplementary material for the interested reader.
4. S OLVING THE S CHR ¨ ODINGER E QUATION
As a ﬁrst experiment to demonstrate the correctness of the method on a problem with a known solution, we investigated the use of SpIN for solving the Schr¨odinger equation for a two-dimensional hydrogen atom.
The time-independent Schr¨odinger equation for a single particle with mass m in a potential ﬁeld V ( x ) is a partial differential equation of the form:
[Equation]
20) whose solutions describe the wavefunctions ψ ( x ) with unique energy E .
The probability of a particle being at position x then has the density | ψ ( x ) | 2 .
The solutions are eigenfunctions of the linear operator H (cid:44) − h 2 2 m ∇
+ V ( x ) — known as the Hamiltonian operator.
We set (cid:126) 2 2 m to 1 and choose
[Figure]
V ( x ) = 1 | x | , which corresponds to the potential from a charged particle.
In 2 or 3 dimensions this can be solved exactly, and in 2 dimensions it can be shown that there are 2 n + 1 eigenfunctions with energy − 1 (2 n +1) 2 for all n = 0 , 1 , 2 , . . .
(Yang et al., 1991).
We trained a standard neural network to approximate the wavefunction ψ ( x ) , where each unit of the output layer was a solution with a different energy E .
Details of the training network and experimental setup are given in the supplementary material in Sec.
C.1.
We found it critical to set the decay rate for RMSProp to be slower than the decay β used for the moving average of the covariance in SpIN, and expect the same would be true for other adaptive gradient methods.
To investigate the effect of biased gradients and demonstrate how SpIN can correct it, we speciﬁcally chose a small batch size for our experiments.
As an additional baseline over the known closed-form solution, we computed eigenvectors of a discrete approximation to H on a 128 × 128 grid.
Training results are shown in Fig.
In Fig.
1a, we see the circular harmonics that make up the electron orbitals of hydrogen in two dimensions.
With a small batch size and no bias correction, the eigenfunctions (Fig. 1b) are incorrect and the eigenvalues (Fig.
1d, ground truth in black) are nowhere near the true minimum.
With the bias correction term in SpIN, we are able to both accurately estimate the shape of the eigenfunctions (Fig. 1c) and converge to the true eigenvalues of the system (Fig. 1e).
Note that, as eigenfunctions 2-4 and 5-9 are nearly degenerate, any linear combination of them is also an eigenfunction, and we do not expect Fig.
1a and Fig.
1c to be identical.
The high accuracy of the learned eigenvalues gives strong empirical support for the correctness of our method.
4. D EEP S LOW F EATURE A NALYSIS
Having demonstrated the effectiveness of SpIN on a problem with a known closed-form solution, we now turn our attention to problems relevant to representation learning in vision.
We trained a convolutional neural network to extract features from videos, using the Slow Feature Analysis
[Figure] [Figure]
[Figure] [Figure]
[Figure]
kernel of Eq.
The video is a simple example with three bouncing balls.
The velocities of the balls are constant until they collide with each other or the walls, meaning the time dynamics are reversible, and hence the transition function is a symmetric operator.
We trained a model with 12 output eigenfunctions using similar decay rates to the experiments in Sec. 4.1.
Full details of the training setup are given in Sec.
C.2, including training curves in Fig.
During the course of training, the order of the different eigenfunctions often switched, as lower eigenfunctions sometimes took longer to ﬁt than higher eigenfunctions.
Analysis of the learned solution is shown in Fig.
Fig.
2a is a heatmap showing whether the feature is likely to be positively activated (red) or negatively activated (blue) when a ball is in a given position.
Since each eigenfunction is invariant to change of sign, the choice of color is arbitrary.
Most of the eigenfunctions are encoding for the position of balls independently, with the ﬁrst two eigenfunctions discovering the separation between up/down and left/right, and higher eigenfunctions encoding higher frequency combinations of the same thing.
However, some eigenfunctions are encoding more complex joint statistics of position.
For instance, one eigenfunction (outlined in green in Fig.
2a) has no clear relationship with the marginal position of a ball.
But when we plot the frames that most positively or negatively activate that feature (Fig. 2b) we see that the feature is encoding whether all the balls are crowded in the lower right corner, or one is there while the other two are far away.
Note that this is a fundamentally nonlinear feature, which could not be discovered by a shallow model.
Higher eigenfunctions would likely encode for even more complex joint relationships.
None of the eigenfunctions we investigated seemed to encode anything meaningful about velocity, likely because collisions cause the velocity to change rapidly, and thus optimizing for slowness of features is unlikely to discover this.
A different choice of kernel may lead to different results.
D ISCUSSION
We have shown that a single uniﬁed framework is able to compute spectral decompositions by stochastic gradient descent on domains relevant to physics and machine learning.
This makes it possible to learn eigenfunctions over very high-dimensional spaces from very large datasets and generalize to new data without the Nystr¨om approximation.
This extends work using slowness as a criterion for unsupervised learning without a generative model, and addresses an unresolved issue with biased gradients due to ﬁnite batch size.
A limitation of the proposed solution is the requirement of computing full Jacobians at every time step, and improving the scaling of training is a promising direction for future research.
The physics application presented here is on a fairly simple system, and we hope that Spectral Inference Nets can be fruitfully applied to more complex physical systems for which computational solutions are not yet available.
The representations learned on video data show nontrivial structure and sensitivity to meaningful properties of the scene.
These representations could be used for many downstream tasks, such as object tracking, gesture recognition, or faster exploration and subgoal discovery in reinforcement learning.
Finally, while the framework presented here is quite general, the examples shown investigated only a small number of linear operators.
Now that the basic framework has been laid out, there is a rich space of possible kernels and architectures to combine and explore.
Supplementary Material for “Spectral Inference Networks” A B REAKING THE S YMMETRY B ETWEEN E IGENFUNCTIONS
Since Eq. 6 is invariant to linear transformation of the features u ( x ) , optimizing it will only give a function that spans the top K eigenfunctions of K .
We discuss some of the possible ways to recover ordered eigenfunctions, explain why we chose the approach of using masked gradients, and provide a derivation of the closed form expression for the masked gradient in Eq.
A. A LTERNATIVE S TRATEGIES TO B REAK S YMMETRY
If u ∗ ( x ) is a function to R N that is an extremum of Eq.
6, then E x (cid:48)
[ k ( x , x (cid:48) )
u ∗ ( x (cid:48) )]
=
Ωu ∗ ( x ) for some matrix Ω ∈ R N × N which is not necessarily diagonal .
We can express this matrix in terms of quantities in the objective:
[Equation] [Equation]
To transform u ∗ ( x ) into ordered eigenfunctions, ﬁrst we can orthogonalize the functions by multiplying by L − 1 where L is the Cholesky decomposition of Σ .
Let v ∗ ( x ) =
L − 1 u ∗ ( x ) , then
E x (cid:48)
[Equation]
x ,
[Equation]
v
[Equation] [Equation]
(cid:48)
[Equation]
The matrix L − 1 ΠL − T = Λ is symmetric, so we can diagonalize it: Λ = VDV T , and then w ∗ ( x ) =
V T v ∗ ( x ) =
V T L − 1 u ∗ ( x ) are true eigenfunctions, with eigenvalues along the diagonal of D .
In principle, we could optimize Eq.
6, accumulating statistics on Π and Σ , and transform the functions u ∗ into w ∗ at the end.
In practice, we found that the extreme eigenfunctions were “contaminated” by small numerical errors in the others eigenfunctions, and that this approach struggled to learn degenerate eigenfunctions.
This inspired us to explore the masked gradient approach instead, which improves numerical robustness.
A. C ORRECTNESS OF M ASKED G RADIENT A PPROACH
Throughout this section, let x i : j be the slice of a vector from row i to j and let A i : j,k : (cid:96) be the block of a matrix A containing rows i through j and columns k through (cid:96) .
Let Σ =
E x
[ u ( x ) u ( x ) T ] , Π = E x , x (cid:48)
[ k ( x , x (cid:48) )
u ( x ) u ( x (cid:48) )
T ] , L be the Cholesky decomposition of Σ and Λ =
L − 1 ΠL − T .
The arguments here are not meant as mathematically rigorous proofs but should give the reader enough of an understanding to be conﬁdent that the numerics of our method are correct for optimization over a sufﬁciently expressive class of functions.
Claim 1.
Λ 1: n, 1: n is independent of u n +1: n ( x ) .
The Cholesky decomposition of a positive-deﬁnite matrix is the unique lower triangular matrix with positive diagonal such that LL T = Σ .
Expanding this out into blocks yields:
[Equation]
Inspecting the upper left block, we see that L 1: n, 1: n L T 1: n, 1: n = Σ 1: n, 1: n .
As L 1: n, 1: n is also lower- triangular, it must be the Cholesky decomposition of Σ 1: n, 1: n .
The inverse of a lower triangular matrix will also be lower triangular, and a similar argument to the one above shows that the upper left block of the inverse of a lower triangular matrix will be the inverse of the upper left block, so the upper left block of Λ can be written as:
[Equation]
which depends only on u 1: n ( x )
Claim 2.
Let ˜ ∇ u Tr( Λ )
= (cid:16) ∂ Λ 11 ∂u 1 , . . .
, ∂ Λ NN ∂u N (cid:17) , and let v ( x ) =
L − 1 u ( x ) .
If the parameters θ of u ( x ) are maximized by gradient ascent so that ˜ ∇ θ Tr( Λ ) = 0 , and the true eigenfunctions are in the class of functions parameterized by θ , then v ( x ) will be the eigenfunctions of the operator K deﬁned as K [ f ]( x ) =
E x (cid:48)
[ k ( x , x (cid:48) )
f ( x (cid:48) )] , ordered from highest eigenvalue to lowest.
The argument proceeds by induction.
Λ 11 =
L − 1 11 Π 11 L − 1
=
Π 11 Σ 11 , which is simply the Rayleigh quotient in Eq.
10 for i = 1 .
The maximum of this is clearly proportional to the top eigenfunction, and v 1 ( x ) =
L − 1 11 u 1 ( x ) =
Σ − 1 / 2u 1 ( x ) is the normalized eigenfunction.
Now suppose v 1: n ( x ) =
L − 1 1: n, 1: n u 1: n ( x ) are the ﬁrst n eigenfunctions of K .
Because u 1: n ( x ) span the ﬁrst n eigenfunctions, and Λ ii is independent of u n +1 ( x ) for i < n +1 , u 1: n +1 ( x ) is a maximum of (cid:80) ni
=1 Λ ii no matter what the function u n +1 ( x ) is.
Training u n +1 ( x ) with the masked gradient ˜ ∇ u Tr( Λ ) is equivalent to maximizing Λ ( n +1)( n +1) , so for the optimal u n +1 ( x ) , u 1: n +1 ( x ) will be a maximum of (cid:80) n +1 i =1 Λ ii .
Therefore u 1: n ( x ) span the ﬁrst n eigenfunctions and u 1: n +1 ( x ) span the ﬁrst n + 1 eigenfunctions, so orthogonalizing u 1: n +1 ( x ) by multiplication by L − 1 1: n +1 , 1: n +1 will subtract anything in the span of the ﬁrst n eigenfunctions off of u n +1 ( x ) , meaning v n +1 ( x ) will be the ( n + 1) th eigenfunction of K
A. D ERIVATION OF M ASKED G RADIENT
The derivative of the normalized features with respect to parameters can be expressed as
[Equation] [Equation]
if we ﬂatten out the matrix-valued L , Σ and Π .
The reverse-mode sensitivities for the matrix inverse and Cholesky decomposition are given by ¯ A = − C T ¯ CC T where
C =
A − 1 and ¯ Σ
=
L − T Φ ( L T ¯ L ) L − 1 where L is the Cholesky decomposition of Σ and Φ ( · ) is the operator that replaces the upper triangular of a matrix with its lower triangular transposed (Giles, 2008; Murray, 2016).
Using this, we can compute the gradients in closed form by application of the chain rule.
To simplify notation slightly, let ∆ k and Φ
k be matrices deﬁned as:
[Equation]
Then the unmasked gradient has the form:
[Equation]
[Equation]
(28) while the gradients of Π and Σ with respect to u are given (elementwise) by:
[Equation]
which, in combination, give the unmasked gradient with respect to u as:
[Equation]
Here the gradient is expressed as a row vector
[Equation]
to be consistent with Eq.
[Equation]
and a factor of 2 has been dropped in the last line that can be absorbed into the learning rate.
To zero out the relevant elements of the gradient ∇ u Λ kk as described in Eq.
[Equation]
we can rightmultiply by ∆ k .
The masked gradients can be expressed in closed form as:
[Equation] [Equation]
35) where triu and diag give the upper triangular and diagonal of a matrix, respectively.
A TensorFlow implementation of this masked gradient is given below.
B T ENSOR F LOW I MPLEMENTATION OF S P IN U PDATE
Here we provide a short pseudocode implementation of the updates in Alg. 1 in TensorFlow.
The code is not intended to run as is, and leaves out some global variables, proper initialization and code for constructing networks and kernels.
However all nontrivial elements of the updates are given in detail here.
import tensorflow as tf from tensorflow.python.ops.parallel_for import jacobian @tf.custom_gradient defcovariance(x,y): batch_size=
float(x.shape[0].value) cov=
tf.matmul(x,y,transpose_a=
True) / batch_size def gradient(grad): return (tf.matmul(y, grad) / batch_size, tf.matmul(x, grad) / batch_size) return cov, gradient @tf.custom_gradient def eigenvalues(sigma, pi): """Eigenvalues as custom op so that we can overload gradients."
"" chol = tf.cholesky(sigma) choli = tf.linalg.inv(chol) rq = tf.matmul(choli, tf.matmul(pi, choli, transpose_b=
True)) eigval = tf.matrix_diag_part(rq)
def gradient(_): """Symbolic form of the masked gradient."
"" dl = tf.diag(tf.matrix_diag_part(choli)) triu = tf.matrix_band_part(tf.matmul(rq, dl), 0, -1) dsigma =
-1.0*tf.matmul(choli, triu, transpose_a=
True) dpi = tf.matmul(choli, dl, transpose_a=
True) return dsigma, dpi return eigval, gradient def moving_average(x, c): """Creates moving average operation.
This is pseudocode for clarity!
Should actually initialize sigma_avg with tf.eye, and should add handling for when x is a list.
""" ma = tf.
Variable(tf.zeros_like(x), trainable=False) ma_update =
tf.assign(ma, (1-c)*ma + c*x) return ma, ma_update def spin(x1, x2, network, kernel, params, optim): """ to create TensorFlow ops for learning in SpIN.
Args: neural parameters
Returns: step:op that implements one iteration of SpIN training update eigenfunctions: op that gives ordered eigenfunctions """ # ‘u1‘ and ‘u2‘ are assumed to have the batch elements along first
# dimension and different eigenfunctions along the second dimension u1 = network(x1, params) u2 = network(x2,
params
) sigma = 0.5 * (covariance(u1, u1) + covariance(u2, u2)) sigma.set_shape((u1.shape[1], u1.shape[1]))
# For the hydrogen examples in Sec. 4.1, ‘kernel(x1, x2)*u2‘ # can be replaced by the result of applying the operator # H to the function defined by ‘network(x1, params)‘.
pi =
[Table]
u1, kernel(
[Table] [Table]
)*u2) pi.
[Table]
_shape((u1.shape[1], u1.shape[1]))
[Table]
with tf.control_dependencies(update_sigma_jac + [update_sigma]): eigval = eigenvalues(sigma_avg, pi) loss = tf.reduce_sum(eigval)
sigma_back = tf.gradients(loss, sigma_avg)[0] gradients =
[ tf.tensordot(sigma_back, sig_jac, [[0, 1], [0, 1]]) + grad for sig_jac, grad in zip(sigma_jac_avg, tf.gradients(loss, params)) ] step = optim.apply_gradients(zip(gradients, params)) eigenfunctions =
tf.matmul(u1, tf.linalg.inv(tf.cholesky(sigma_avg)), transpose_b=
True) return step, eigenfunctions
C E XPERIMENTAL D ETAILS C. S OLVING THE S CHR ¨ ODINGER E QUATION
To solve for the eigenfunctions with lowest eigenvalues, we used a neural network with 2 inputs (for the position of the particle), 4 hidden layers each with 128 units, and 9 outputs, corresponding to the ﬁrst 9 eigenfunctions.
We used a batch size of 128 - much smaller than the 16,384 nodes in the 2D grid used for the exact eigensolver solution.
We chose a softplus nonlinearity log(1+exp( x )) rather than the more common ReLU, as the Laplacian operator ∇ 2 would be zero almost everywhere for a ReLU network.
We used RMSProp (Tieleman & Hinton, 2012) with a decay rate of 0.999 and learning rate of 1e-5 for all experiments.
We sampled points uniformly at random from the box [ − D, D ] 2 during training, and to prevent degenerate solutions due to the boundary condition, we multiplied the output of the network by (cid:81) i ( (cid:112) 2 D 2 − x 2 i − D ) , which forces the network output to be zero at the boundary without the derivative of the output blowing up.
We chose D = 50 for the experiments shown here.
We use the ﬁnite difference approximation of the differential Laplacian given in Sec.
3.2 with (cid:15) some small number (around 0.1), which takes the form:
[Equation]
when applied to ψ ( x ) .
Because the Hamiltonian operator is a purely local function of ψ ( x ) , we don’t need to sample pairs of points x , x (cid:48) for each minibatch, which simpliﬁes calculations.
We made one additional modiﬁcation to the neural network architecture to help separation of different eigenfunctions.
Each layer had a block-sparse structure that became progressively more separated the deeper into the network it was.
For layer (cid:96) out of L with m inputs and n outputs, the weight w ij was only nonzero if there exists k ∈ { 1 , . .
, K } such that i ∈
[ k − 1 K − 1 (cid:96) − 1 L m, k − 1 K − 1 (cid:96) − 1 L m +
L − (cid:96) +
L m ] and j ∈ [ k − 1 K − 1 (cid:96) − 1 L n, k − 1 K − 1 (cid:96) − 1 L n + L − (cid:96) +
L n ] .
This split the weight matrices into overlapping blocks, one for each eigenfunction, allowing features to be shared between eigenfunctions in lower layers of the network while separating out features which were distinct between eigenfunctions higher in the network.
C. D EEP S LOW F EATURE A NALYSIS
We trained on 200,000 64 × 64 pixel frames, and used a network with 3 convolutional layers, each with 32 channels, 5 × 5 kernels and stride 2, and a single fully-connected layer with 128 units before outputting 12 eigenfunctions.
We also added a constant ﬁrst eigenfunction, since the ﬁrst eigenfunction of the Laplacian operator is always constant with eigenvalue zero.
This is equivalent to forcing the features to be zero-mean.
We used the same block-sparse structure for the weights that was used in the Schr¨odinger equation experiments, with sparsity in weights between units extended to sparsity in weights between entire feature maps for the convolutional layers.
We trained with RMSProp with learning rate 1e-6 and decay 0.999 and covariance decay rate β = 0 .
01 for 1,000,000 iterations.
To make the connection to gradient descent clearer, we use the opposite convention to RMSProp: β = 1 corresponds to zero memory for the moving average, meaning the RMS term in RMSProp decays ten times more slowly than the covariance moving average in these experiments.
Each batch contained 24 clips of 10 consecutive frames.
So that the true state was fully observable, we used two consecutive frames as the input x t , x t +1 and trained the network so that the difference from that and the features for the frames x t +1 , x t +2 were as small as possible.
C. S UCCESSOR F EATURES AND THE A RCADE L EARNING E NVIRONMENT
We provide a qualitative comparison of the performance of SpIN with the SFA objective against the successor feature approach for learning eigenpurposes Machado et al.
(2018) on the Arcade Learning
Environment (Bellemare et al., 2013).
As in Machado et al.
(2018), we trained a network to perform next-frame prediction on 500k frames of a random agent playing one game.
We simultaneously trained another network to compute the successor features (Barreto et al., 2017) of the latent code of the next-frame predictor, and computed the “eigenpurposes” by applying PCA to the successor features on 64k held-out frames of gameplay.
We used the same convolutional network architecture as Machado et al.
(2018), a batch size of 32 and RMSProp with a learning rate of 1e-4 for 300k iterations, and updated the target network every 10k iterations.
While the original paper did not mean-center the successor features when computing eigenpurposes, we found that the results were signiﬁcantly improved by doing so.
Thus the baseline presented here is actually stronger than in the original publication.
On the same data, we trained a spectral inference network with the same architecture as the encoder of the successor feature network, except for the fully connected layers, which had 128 hidden units and 5 non-constant eigenfunctions.
We tested SpIN on the same 64k held-out frames as those used to estimate the eigenpurposes.
We used the same training parameters and kernel as in Sec.
As
SpIN is not a generative model, we must ﬁnd another way to compare the features learned by each method.
We averaged together the 100 frames from the test set that have the largest magnitude positive or negative activation for each eigenfunction/eigenpurpose.
Results are shown in Fig.
4, with more examples and comparison against PCA on pixels at the end of this section.
By comparing the top row to the bottom row in each image, we can judge whether that feature is encoding anything nontrivial.
If the top and bottom row are noticeably different, this is a good indication that something is being learned.
It can be seen that for many games, successor features may ﬁnd a few eigenpurposes that encode interesting features, but many eigenpurposes do not seem to encode anything that can be distinguished from the mean image.
Whereas for SpIN, nearly all eigenfunctions are encoding features such as the presence/absence of a sprite, or different arrangements of sprites, that lead to a clear distinction between the top and bottom row.
Moreover, SpIN is able to learn to encode these features in a fully end-to-end fashion, without any pixel reconstruction loss, whereas the successor features must be trained from two distinct losses, followed by a third step of computing eigenpurposes.
The natural next step is to investigate how useful these features are for exploration, for instance by learning options which treat these features as rewards, and see if true reward can be accumulated faster than by random exploration.
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]