Question: Represent the scientific article to find a short TLDR

Question: Variational inference (VI) methods and especially variational autoencoders (VAEs) specify scalable generative models that enjoy an intuitive connection to manifold learning --- with many default priors the posterior/likelihood pair $q(z|x)$/$p(x|z)$ can be viewed as an approximate homeomorphism (and its inverse) between the data manifold and a latent Euclidean space. However, these approximations are well-documented to become degenerate in training. Unless the subjective prior is carefully chosen, the topologies of the prior and data distributions often will not match. Conversely, diffusion maps (DM) automatically \textit{infer} the data topology and enjoy a rigorous connection to manifold learning, but do not scale easily or provide the inverse homeomorphism. In this paper, we propose \textbf{a)} a principled measure for recognizing the mismatch between data and latent distributions and \textbf{b)} a method that combines the advantages of variational inference and diffusion maps to learn a homeomorphic generative model. The measure, the \textit{locally bi-Lipschitz property}, is a sufficient condition for a homeomorphism and easy to compute and interpret. The method, the \textit{variational diffusion autoencoder} (VDAE), is a novel generative algorithm that first infers the topology of the data distribution, then models a diffusion random walk over the data. To achieve efficient computation in VDAEs, we use stochastic versions of both variational inference and manifold learning optimization. We prove approximation theoretic for the dimension dependence of VDAEs, and that locally isotropic sampling in the latent space in a random walk over the reconstructed manifold. Finally, we demonstrate the utility of our method on various real and synthetic datasets, and show that it exhibits performance superior to other generative models. Recent developments in generative models such as variational auto-encoders and generative adversarial networks have made it possible to sample remarkably realistic points from complex high dimensional distributions at low computational cost. While their methods are very different -one is derived from variational inference and the other from game theory -their ends both involve learning smooth mappings from a user-defined prior distribution to the modeled distribution. These maps are closely tied to manifold learning when the prior is supported over a Euclidean space (e.g. Gaussian or uniform priors) and the data lie on a manifold (also known as the Manifold Hypothesis, see ;). This is because manifolds themselves are defined by sets that have homeomorphisms to such spaces. Learning such maps is beneficial to any machine learning task, and may shed light on the success of VAEs and GANs in modeling complex distributions. Furthermore, the connection to manifold learning may explain why these generative models fail when they do. Known as posterior collapse in VAEs (; ; ;) and mode collapse in GANs , both describe cases where the forward/reverse mapping to/from Euclidean space collapses large parts of the input to a single output. This violates the bijective requirement of the homeomorphic mapping. It also in degenerate latent spaces and poor generative performance. A major cause of such failings is when Figure 1: A diagram depicting one step of the diffusion process modeled by the variational diffusion autoencoder (VDAE). The diffusion and inverse diffusion maps ψ, ψ −1, as well as the covariance C of the random walk on M Z, are all approximated by neural networks. the geometries of the prior and target data do not agree. We explore this issue of prior mismatch and previous treatments of it in Section 3. Given their connection to manifold learning, it is natural to look to classical approaches in the field for ways to improve VAEs. One of the most principled methods is spectral learning (Schölkopf et al., 1998; ;) which involves describing data from a manifold X ⊂ M X by the eigenfunctions of a kernel on M X. We focus specifically on DMs, where show that normalizations of the kernel approximate a very specific diffusion process, the heat kernel over M X. A crucial property of the heat kernel is that, like its physical analogue, it defines a diffusion process that has a uniform stationary distribution -in other words, drawing from this stationary distribution draws uniformly from the data manifold. established another crucial property of DMs, namely that distances in local neighborhoods in the eigenfunction space are nearly isometric to corresponding geodesic distances on the manifold. However, despite its strong theoretical guarantees, DMs are poorly equipped for large scale generative modeling as they are not easily scalable and do not provide an inverse mapping from the intrinsic feature space. In this paper we address issues in variational inference and manifold learning by combining ideas from both. Theory in manifold learning allows us to better recognize prior mismatch, whereas variational inference provides a method to learn the difficult to approximate inverse diffusion map. Our contributions: 1) We introduce the locally bi-Lipschitz property, a sufficient condition for a homeomorphism, for measuring the stability of a mapping between latent and data distributions. 2) We introduce VDAEs, a class of variational autoencoders whose encoder-decoder feedforward pass approximates the diffusion process on the data manifold with respect to a user-defined kernel k. 3) We show that deep neural networks are capable of learning such diffusion processes, and 4) that networks approximating this process produce random walks that have certain desirable properties, including well defined transition and stationary distributions. 5) Finally, we demonstrate the utility of the VDAE framework on a set of real and synthetic datasets, and show that they have superior performance and satisfy the locally bi-Lipschitz property where GANs and VAEs do not. Variational inference is a machine learning method that combines Bayesian statistics and latent variable models to approximate some probability density p(x). VI assumes and exploits a latent variable structure in the assumed data generation process, that the observations x ∼ p(x) are conditionally distributed given unobserved latent vari-ables z. By modeling the conditional distribution, then marginalizing over z, as in we obtain the model evidence, or likelihood that x could have instead been drawn from p θ (x). Maximizing Eq. 1 leads to an algorithm for finding likely approximations of p(x). As the cost of computing this integral scales exponentially with the dimension of z, we instead maximize the evidence lower bound (ELBO): where q(z|x) is usually an approximation of p θ (z|x). Optimizing the ELBO is sped up by taking stochastic gradients , and further accelerated by learning a global function approximator q φ in an autoencoding structure . Diffusion maps on the other hand, are a class of kernel methods that perform non-linear dimensionality reduction on a set of observations X ⊆ M X, where M X is the data manifold. Given a symmetric and positive kernel k, DM considers the induced random walk on the graph of X, where given x, y ∈ X, the transition probabilities p(y|x) = p(x, y) are row normalized versions of k(x, y). Moreover, the diffusion map ψ embeds the data X ∈ R m into the Euclidean space R D so that the diffusion distance is approximated by Euclidean distance. This is a powerful property, as it allows the arbitrarily complex random walk induced by k on M X to become an isotropic Gaussian random walk on ψ(M X). SpectralNet is an algorithm introduced by algorithm in Shaham et al. (2018b) to speed up the diffusion map. Until recently, the method ψ k could only be computed via the eigendecomposition of K. As a , DMs were only be tractable for small datasets, or on larger datasets by combining landmark-based estimates and Nystrom approximation techniques. However, Shaham et al. (2018b) propose approximations of the function ψ itself in the case that the kernel k is symmetric. In particular, we will leverage SpectralNet to enforce our diffusion embedding prior. Locally bi-lipschitz coordinates by kernel eigenfunctions. analyzed the construction of local coordinates of Riemannian manifolds by Laplacian eigenfunctions and diffusion map coordinates. They establish, for all x ∈ X, the existence of some neighborhood U (x) and d spectral coordinates given U (x) that define a bi-Lipschitz mapping from U (x) to R d. With a smooth compact Riemannian manifold, U (x) can be chosen to be a geodesic ball with radius a constant multiple of the inradius (the radius of the largest possible ball around x without intersecting with the manifold boundary), where the constant is uniform for all x, but the indices of the d spectral coordinates as well as the local bi-Lipschitz constants may depend on x. Specifically, the Lipschitz constants involve inverse of the inradius at x multiplied again by some global constants. For completeness we give a simplified statement of the in the supplementary material. Using the compactness of the manifold, one can always cover the manifold with m many neighborhoods (geodesic balls) on which the bi-Lipschitz property in holds. As a , there are a total of D spectral coordinates, D ≤ md (in practice D is much smaller than md, since the selected spectral coordinates in the proof of tend to be low-frequency ones, and thus the selection on different neighborhoods tend to overlap), such that on each of the m neighborhoods, there exists a subset of d spectral coordinates out of the D ones which are bi-Lipschitz on the neighborhood, and the Lipschitz constants can be bounded uniformly from below and above. Our proposed measure and model is motivated by degenerate latent spaces and poor generative performance in a variational inference framework arising from prior mismatch: when the topologies of the data and prior distributions do not agree. In real world data, this is usually due to two factors: first, when the dimensionalities of the distributions do not match, and second, when the geometries do not match. It is easy to see that homeomorphisms between the distributions will not exist in either case: pointwise correspondences cannot be established, thus the bijective condition cannot be met. As a , the model has poor generative performance -for each point not captured in the pointwise correspondence, the latent or generated distribution loses expressivity. Though the default choice of Gaussian distribution for p(z) is mathematically elegant and computationally expedient, there are many datasets, real and synthetic, for which this distribution is ill-suited. It is well known that spherical distributions are superior for modeling directional data , which can be found in fields as diverse as bioinformatics , geology , material science , natural image processing , and simply preprocessed datasets 1. Additionally observe that no homeomorphism exists between R k and S 1 for any k. For data distributed on more complex manifolds, the literature is sparse due to the difficult nature of such study. However, the manifold hypothesis is well-known and studied . Previous research on alleviating prior mismatch exists.; consider VAEs with the von-Mises Fisher prior, a geometrically hyperspherical prior. further model arbitrarily complex manifolds as priors, but require explicit knowledge of the manifold (i.e. its projection map, scalar curvature, and volume). consider mixtures of any pre-existing priors. But while these methods increase the expressivity of the priors available, they do not prescribe a method for choosing the prior itself. That responsibility still lies with the user. Convserly, our method chooses the best prior automatically. To our knowledge, ours is the first to take a data-driven approach to prior selection. By using some data to inform the prior, we not only guarantee the existence of a homeomorphism between data and prior distributions, we explicitly define it by the learned diffusion mapψ. In this section we propose VDAEs, a variational inference method that, given the data manifold M X, observations X ⊂ M X, and a kernel k, models the geometry of X by approximating a random walk over the latent diffusion manifold M Z:= ψ(M X). The model is trained by maximizing the local evidence: the evidence (i.e. log-likelihood) of each point given its random walk neighborhood. Points are generated from the trained model by sampling from π, the stationary distribution of the ing random walk. Starting from some point x ∈ X, we can roughly describe one step of the walk as the composition of three functions: 1) the approximate diffusion mapψ Θ: M X → M Z, 2) a sampling procedure from the learned diffusion process z ∼ q φ (z |x) = N (ψ Θ (x),C φ ) on M Z, and 3) the learned inverse diffusion mapψ where the constant c is user-defined and fixed. We rely crucially on three advantages of our latent spaceψ Θ (X): a) that it is well-defined (given the first D eigenvalues of k are distinct), b) well-approximated (given SpectralNet) and c) that Euclidean distances in M Z approximate single-step random walk distances on M X (see Section 2 and). Thus the transition probabilities induced by k can be approximated by Gaussian kernels 2 in M Z. Therefore, to model a diffusion random walk over M Z, we must learn the functionsψ Θ,ψ −1 θ,C φ that approximate the diffusion map, its inverse, and the covariance of the random walk on M Z, at all points z ∈ M Z. SpectralNet gives usψ Θ. To learnψ −1 θ andC φ, we use variational inference. Formally, let us define U x:= B d (x, δ) ∩ M X, where B d (x, δ) as the δ-ball around x with respect to d(·, ·), the diffusion distance on M Z. For each x ∈ X, we define the local evidence of x as where p(x |x)| Ux is the restriction of p(x |x) to U x. The ing local evidence lower bound is: Note that the neighborhood reconstruction error should be differentiated from the self reconstruction error that is in VAEs. Eq. 4 produces the empirical loss function: where is the deterministic, differentiable function, depending onψ Θ andC φ, that generates q φ by the reparameterization trick 3 . Algorithm 1 VDAE training Θ, φ, θ ← Initialize parameters Obtain parameters Θ for the approximate diffusion mapψ Θ by Shaham et al. (2018b) while not converged do Take one step of the diffusion random walk Compute gradients of the loss, i.e. Eq. equation 4 Update φ, θ using g Here we discuss the algorithm for generating data points from p(x). Composing q φ (z |x)(≈ p θ (z |x)) with p θ (x |z) gives us an approximation of p θ (x |x). Then the simple, parallelizable, and fast random walk based sampling procedure naturally arises: initialize with an arbitrary point on the manifold x 0 ∈ M X, then pick suitably large N and for n = 1,..., N draw x n ∼ p(x|x n−1). Eventually, our diffusion random walk converges on its stationary distribution π. , this is guaranteed to be the uniform distribution on the data manifold. See Section 6.2 for examples of points drawn from this procedure. We now introduce a practical implementation VDAEs, considering the case whereψ Θ (x), q φ (z |x) and p θ (x |z) are neural network functions, as they are in VAEs and SpectralNet, respectively. The neighborhood reconstruction error. Since q φ (z |x) models the neighborhood ofψ Θ (x), we may sample q φ to obtain z (the neighbor of x in the latent space). This gives p θ (x |x) ≈ ψ −1 θ (q φ (z |x)), where ψ −1 exists due to the bi-Lipschitz property. We can efficiently approximate x ∈ M X by considering the closest embedded data pointψ Θ (x) ∈ M Z to z =ψ Θ (x). This is because Euclidean distance on M Z approximates the diffusion distance on M X. In other words, x ∼ p θ (x |x) ≈ψ −1 θ (q φ (z |x)) which we approximate empirically by where A ⊆ X is the training batch. On the other hand, the divergence of random walk distributions −D KL (q φ (z |x)||p θ (z |x)) can be modeled simply as the divergence of two Gaussian kernels defined on M Z. Though p θ (z |x) is intractable, the diffusion map ψ gives us the diffusion embedding Z, which is an approximation of the true distribution of p θ (z |x) in a neighborhood around z = ψ(x). We estimate the first and second moments of this distribution in R D by computing the local Mahalanobis distance of points in the neighborhood. Then, by minimizing the KL divergence between q φ (z |x) and the one implied by this Mahalanobis distance, we obtain the loss: where is the covariance of the points in a neighborhood of z = ψ(x) ∈ Z, and α is a scaling parameter. Note that C φ (x) does not have to be diagonal, and in fact is most likely not. Combining Eqs. 6 and 7 we obtain Algorithm 1. Now we consider the sampling procedure. Since we use neural networks to approximate q φ (z |x) and p θ (x |z), the generation procedure is highly parallelizable. We empirically observe the random walk enjoys rapid mixing properties -it does not take many iterations of the random walk to sample from all of M Z 4. This leads to Algorithm 2. Algorithm 2 VDAE sampling Take one step of the diffusion random walk Map back into input space t ← t + 1 We theoretically prove that the desired inverse map ψ −1 from spectral coordinate codes back to the manifold can be approximated by a decoder network, where the network complexity is bounded by quantities related to the intrinsic geometry of the manifold. This section relies heavily on the known bi-Lipschitz property of , which we are approximating with the VDAE latent space without the need for regularization. The theory for the capacity of the encoder to map M to the diffusion map space ψ(M) has already been considered in Shaham et al. (2018a) and. We instead focus on the decoder, which requires a different treatment. The following theorem is proved in Appendix A.3, based upon the in. to have a subset of coordinates that are locally bi-Lipschitz. Let X = [X 1, ..., X m] be the set of all m extrinsic coordinates of the manifold. Then there exists a sparsely-connected ReLU network f N, with 4DC M X nodes in the first layer, 8dmN nodes in the second layer, and 2mN nodes in the third layer, and m nodes in the output layer, such that where the norm is interpreted as Here C ψ depends on how sparsely X(ψ(x)) Ui can be represented in terms of the ReLU wavelet frame on each neighborhood U i, and C M X on the curvature and dimension of the manifold M X. Theorem 1 is complementary to the theorem in Shaham et al. (2018a), which provides guarantees for the encoder, as Theorem 1 demonstrates a similar approximation theoretic argument for the decoder. The proof is built on two properties of ReLU neural networks: 1) their ability to split curved domains into small, almost Euclidean patches, 2) their ability to build differences of bump functions VDAE SVAE VAE GAN Figure 2: Reconstructed images from the rotating bulldog example plotted in the latent space of VDAE (left), Spherical VAE (SVAE, left-middle) and VAE (right-middle), and GAN (right) on each patch, which allows one to borrow approximation from the theory of wavelets on spaces of homogeneous type. The proof also crucially uses the bi-Lipschitz property of the diffusion embedding. The key insight of Theorem 1 is that, because of the bi-Lipschitz property, the coordinates of the manifold in the ambient space R m can be thought of as functions of the diffusion coordinates. We show that because each of coordinates function X i is a Lipschitz function, the ReLU wavelet coefficients of X i are necessarily 1. This allows us to use the existing guarantees of Shaham et al. (2018a) to complete the desired bound. We also discuss the connections between the distribution at each point in diffusion map space, q φ (z|x), and the of this distribution after being decoded through the decoder network f N (z) for z ∼ q φ (z|X). Similar to , we characterize the covariance matrix The following theorem is proved in Appendix A.3. Theorem 2. Let f N be a neural network approximation to X as in Theorem 1, such that it approximates the extrinsic manifold coordinates. Let C ∈ R m×m be the covariance matrix, Σ) with small enough Σ that there exists a patch U z0 ⊂ M around z 0 satisfying the bi-Lipschitz property of , and such that P r(z ∼ q φ (z|x) ∈ ψ(U z0)) <. Then the number of eigenvalues of C greater than is at most d, and C = J z0 ΣJ 6 EXPERIMENTAL We consider the problem of generating new frames from a video of rigid movement. We take 200 frames of a color video (each frame is 100 × 80 × 3) of a spinning bulldog. Due to the spinning of figure and the fixed , this creates a low-dimensional approximately circular manifold. We compare our method to VAE, the (with a bi-lipchitz constraint on the critic), and the hyperspherical. For the VAE, we use a two dimensional Gaussian prior p θ (z), such that z ∼ N (0, I 2). The noise injected to the GAN is drawn from a two dimensional uniform distribution p θ (z), such that z i ∼ U, i = 1, 2. For the spherical VAE, we use a latent dimension of D = 2, which highlights the dimension mismatch issue that occurs with a spherical prior. This is a benefit of VDAE, even if we choose D > d the latent embedding will still only be locally d dimensional. We use the same architecture for all networks which consists of one hidden layer with 512 neurons, activation function for all networks are tanh. In Fig. 2, we present 300 generated samples, by displaying them on a scatter plot with coordinates corresponding to their latent dimensions z 1 and z 2. In this series of experiments, we visualize the of the sampling procedure in Algorithm 2 on three synthetic manifolds. As discussed in 4.2, we randomly select an initial seed point, then recursively sample from p θ (x |x) many times to simulate a random walk on the manifold. In the top row of Fig. 3, we highlight the location of the initial seed point, take 20 steps of the random walk, and display the ing generated points on three learned manifolds. Clearly after a large number of resampling iterations, the algorithm continues to generate points on the manifold, and the distribution of sampled points converges to a uniform stationary distribution on the manifold. Moreover, this stationary distribution is reached very quickly. In the bottom row of the same Fig. 3, we show p θ (x |x) by sampling a large number of points sampled from the single seed point. As can be seen, a single step of p θ (x |x) covers a large part of the latent space. The architecture also uses one hidden layer of 512 neurons and tanh activations. In this section, we deal with the problem of generating samples from data with multiple clusters in an unsupervised fashion (i.e. no a priori knowledge of the cluster structure). Clustered data creates a problem for many generative models, as the topology of the latent space (i.e. normal distribution) differs from the topology of the data space with multiple clusters. In our first experiment, we show that our method is capable of generating new points from a particular cluster given an input point from that cluster. This generation is done in an unsupervised fashion, which is a different setting from the approach of conditional that require training labels. We demonstrate this property on MNIST in Figure 4, and show that the newly generated points after short diffusion time remain in the equivalent class to the seeded image. Here the architecture is a standard fully convolutional architecture. Details can be found in Appendix A.4. Figure 4: An example of cluster conditional sampling with our method, given a seed point (top left of each image grid). The DVAE is able to produce examples via the random walk that stay approximately within the cluster of the seed point, without any supervised knowledge of the cluster. The problem of addressing difference in topologies between the latent space of a generative model and the output data has been acknowledged in recent works about rejection sampling . Rejection sampling of neural networks consists of generating a large collection of samples using a standard GAN, and then designing a probabilistic algorithm to decide in a post-hoc fashion whether the points were truly in the support of the data distribution p(x). In the following experiment, we compare to the standard example in the generative model literature. The data consists of nine bounded spherical densities with significant minimal separation, lying on a 5 × 5 grid. A standard GAN or VAE struggles to avoid generating points in the gaps between Figure 5: Comparison between GAN, DRS-GAN, and our samples on a 5 × 5 Gaussian grid. GAN and DRS-GAN samples taken from. Shown from left-right are Original, GAN, DRS-GAN, and our method. these densities, and thus requires the post-sampling rejection analysis. On the other hand, our model creates a latent space that separates each of these clusters into their own features and only generates points that exist in the neighborhood of training data. Figure 5 clearly shows that this in significantly fewer points generated in the gaps between clusters, as well as eliminating the need to generate additional points that are not in final generated set. Our VDAE architecture here uses one hidden layer of 512 neurons and tanh activations. GAN and DRS-GAN architectures are as described in. Here we describe a practical method for computing the local bi-Lipschitz property, then use it to evaluate several methods on the MNIST dataset. Let Z and X be metric spaces and f: Z → X. We define, for each z ∈ Z and k ∈ N, the function bilip k (z): where Z:= f −1 (X) is the latent embedding of our dataset X 5, d X and d Z are metrics on X and Z, and U z,k is the k-nearest neighborhood of z. Intuitively, increasing values of K can be thought of as an increasing tendency of the learned map to stretch or compress regions of space. By analyzing various statistics of the local bi-Lipschitz measure evaluated at all points of a latent space Z, we can gain insight into how well-behaved a homeomorphism f is. In Table 1 we report the mean and standard deviation, over 10 runs, of the local bi-Lipschitz property for several methods trained on the MNIST dataset. The comparison is between the Wassertein GAN (WGAN), the VAE, the hyperspherical VAE (SVAE), and our method. We use standard architectures prescribed by their respective papers to train the methods. For our method we use a single 500 unit hidden layer network architecture with ReLU nonlinearities for both the encoder and decoder. By constraining our latent space to be the diffusion embedding of the data, our method finds a mapping that automatically enjoys the homeomorphic properties of an ideal mapping, and this is reflected in the low values of the local bi-Lipschitz constant. Conversely, other methods do not consider the topology of the data in the prior distribution. This is especially appparent in the VAE and SVAE, which must generate from the entirety of the input distribution X since they minimize a reconstruction loss. Interestingly, the mode collapse tendency of GANs alleviate the pathology of the bi-Lipschitz constant by allowing the GAN to focus on a subset of the distribution -but this comes at the cost, of course, of collapsing to a few modes of the dataset. Our method is able to reconstruct the entirety of X while simultaneously maintaining a low local bi-Lipschitz constant. We begin with taking the log of the random walk transition likelihood, where q(z) is an arbitrary distribution. We let q(z) to be the conditional distribution q(z |x). Furthermore, if we make the simplifying assumption that p θ (x |z, z) = p θ (x |z), then we obtain Eq. 4 To state the in , we need the following set-up: (C1) M is a d-dimensional smooth compact manifold, possibly having boundary, equipped with a smooth (at least C 2) Riemannian metric g; We denote the geodesic distance by d M, and the geodesic ball centering at x with radius r by B M (x, r). Under (C1), for each point x ∈ M, there exists r M (x) which is the inradius, that is, r is the largest number s.t. B M (x, r) is contained M. Let M be the Laplacian-Beltrami operator on M with Neumann boundary condition, which is self-adjoint on L 2 (M, µ), µ being the Riemannian volume given by g. Suppose that M is re-scaled to have volume 1. The next condition we need concerns the spectrum of the manifold Laplacian (C2) M has discrete spectrum, and the eigenvalues λ 0 ≤ λ 1 ≤ · · · satisfy the Weyl's estimate, i.e. exists constant C which only depends on M s.t. Let ψ j be the eigenfunction associated with λ j, {ψ j} j form an orthonormal bases of L 2 (M, µ). The last condition is (C3) The heat kernel (defined by the heat equation on M) has the spectral representation as That is, Ψ is bi-Lipschitz on the neighborhood B(x, c 1 r M (x)) with the Lipschitz constants indicated as above. The subscript x in Ψ x emphasizes that the indices j 1, · · ·, j d may depend on x. Proof of Theorem 1. The proof of Theorem 1 is actually a simple extension of the following theorem, Theorem 4, which needs to be proved for each individual extrinsic coordinate X k, hence the additional factor of m coming from the L2 norm of m functions. Theorem 4. Let M ⊂ R m be a smooth d-dimensional manifold, ψ(M) ⊂ R D be the diffusion map for D ≥ d large enough to have a subset of coordinates that are locally bi-Lipschitz. Let one of the m extrinsic coordinates of the manifold be denoted X(ψ(x)) for x ∈ M. Then there exists a sparsely-connected ReLU network f N, with 4DC M nodes in the first layer, 8dN nodes in the second layer, and 2N nodes in the third layer, such that where C ψ depends on how sparsely X(ψ(x)) Ui can be represented in terms of the ReLU wavelet frame on each neighborhood U i, and C M on the curvature and dimension of the manifold M. Proof of Theorem 4. The proof borrows from the main theorem of Shaham et al. (2018a). We adopt this notation and summarize the changes in the proof here. For a full description of the theory and guarantees for neural networks on manifolds, see Shaham et al. (2018a). Let C M be the number of neighborhoods U i = B(x i, δ) ∩ M needed to cover M such that ∀x, y ∈ U i, (1 First, we note that as in Shaham et al. (2018a), the first layer of a neural network is capable of using 4D units to select the subset of d coordinates ψ(x) from ψ(x) for x ∈ U i and zeroing out the other D−d coordinates with ReLU bump functions. Then we can define X(ψ(x)) = X(ψ(x)) on x ∈ U i. Now to apply the theorem from Shaham et al. (2018a), we must establish that X Ui: ψ(U i) → R can be written efficiently in terms of ReLU functions. Because of the manifold and diffusion metrics being bi-Lipschitz, we know at a minimum that ψ is invertible on ψ(U i). Because of this invertibility, we will slightly abuse notation and refer to X(ψ(x)) = X(x), where this is understood to be the extrinsic coordinate of the manifold at the point x that cooresponds to ψ(x). we also know that ∀x, y ∈ U i, |X(ψ(x)) − X(ψ(y))| = |X(x) − X(y)| ≤ max where ∇X(z) is understood to be the gradient of X(z) at the point z ∈ M. This means X(ψ(x)) is a Lipschitz function w.r.t. ψ(x). Because X(ψ(x)) Lipschitz continuous, it can be approximated by step functions on a ball of radius 2 − to an error that is at most with the fact that ψ(U i) is compact, gives the fact that on ψ(U i), set of ReLU wavelet coefficients is in 1. And from Shaham et al. (2018a), if on a local patch the function is expressible in terms of ReLU wavelet coefficients in 1, then there is an approximation rate of 1 √ N for N ReLU wavelet terms.

We combine variational inference and manifold learning (specifically VAEs and diffusion maps) to build a generative model based on a diffusion random walk on a data manifold; we generate samples by drawing from the walk's stationary distribution.

Represent the short summary for finding the corresponding scientific article (long):


Question: represent the scientific article to find a short tldr

Question: While deep learning and deep reinforcement learning systems have demonstrated impressive in domains such as image classification, game playing, and robotic control, data efficiency remains a major challenge, particularly as these algorithms learn individual tasks from scratch. Multi-task learning has emerged as a promising approach for sharing structure across multiple tasks to enable more efficient learning. However, the multi-task setting presents a number of optimization challenges, making it difficult to realize large efficiency gains compared to learning tasks independently. The reasons why multi-task learning is so challenging compared to single task learning are not fully understood. Motivated by the insight that gradient interference causes optimization challenges, we develop a simple and general approach for avoiding interference between gradients from different tasks, by altering the gradients through a technique we refer to as “gradient surgery”. We propose a form of gradient surgery that projects the gradient of a task onto the normal plane of the gradient of any other task that has a conflicting gradient. On a series of challenging multi-task supervised and multi-task reinforcement learning problems, we find that this approach leads to substantial gains in efficiency and performance. Further, it can be effectively combined with previously-proposed multi-task architectures for enhanced performance in a model-agnostic way. While deep learning and deep reinforcement learning (RL) have shown considerable promise in enabling systems to perform complex tasks, the data requirements of current methods make it difficult to learn a breadth of capabilities particularly when all tasks are learned individually from scratch. A natural approach to such multi-task learning problems is to train a single network on all tasks jointly, with the aim of discovering shared structure across the tasks in a way that achieves greater efficiency and performance than solving the tasks individually. However, learning multiple tasks all at once in a difficult optimization problem, sometimes leading to worse overall performance and data efficiency compared to learning tasks individually (; a). These optimization challenges are so prevalent that multiple multi-task RL algorithms have considered using independent training as a subroutine of the algorithm before distilling the independent models into a multi-tasking model; a;; ), producing a multi-task model but losing out on the efficiency gains over independent training. If we could tackle the optimization challenges of multi-task learning effectively, we may be able to actually realize the hypothesized benefits of multi-task learning without the cost in final performance. While there has been a significant amount of research in multi-task learning , the optimization challenges are not well understood. Prior work has described varying learning speeds of different tasks and plateaus in the optimization landscape as potential causes, while a range of other works have focused on the model architecture (b;). In this work, we instead hypothesize that the central optimization issue in multi-task learning arises from gradients from different tasks conflicting with one another. In particular, we define two gradients to be conflicting if they point away from one another (i.e., have a negative cosine similarity). As a concrete example, consider the 2D optimization landscapes of two task objectives shown in Figure 1. The optimization landscape of each task consists of a deep valley, as has been characterized of neural network optimization landscapes in the past . When considering the combined optimization landscape for multiple tasks, SGD produces gradients that struggle to efficiently find the optimum. This occurs due to a gradient thrashing phenomenon, where the gradient of one task destabilizes optimization in the valley. We can observe this in Figure 1 (d) when the optimization reaches the deep valley of task 1, but is prevented from traversing the valley to an optimum. In Section 6.2, we find experimentally that this thrashing phenomenon also occurs in a neural network multi-task learning problem. The core contribution of this work is a method for mitigating gradient interference by altering the gradients directly, i.e. by performing "gradient surgery". If two gradients are conflicting, we alter the gradients by projecting each onto the normal plane of the other, preventing the interfering components of the gradient from being applied to the network. We refer to this particular form of gradient surgery as projecting conflicting gradients (PCGrad). PCGrad is model-agnostic, requiring only a single modification to the application of gradients. Hence, it is easy to apply to a range of problem settings, including multi-task supervised learning and multi-task reinforcement learning, and can also be readily combined with other multi-task learning approaches, such as those that modify the architecture. We evaluate PCGrad on multi-task CIFAR classification, multi-objective scene understanding, a challenging multi-task RL domain, and goal-conditioned RL. Across the board, we find PCGrad leads to significant improvements in terms of data efficiency, optimization speed, and final performance compared to prior approaches. Further, on multi-task supervised learning tasks, PCGrad can be successfully combined with prior state-of-the-art methods for multi-task learning for even greater performance. The goal of multi-task learning is to find parameters θ of a model f θ that achieve high average performance across all the training tasks drawn from a distribution of tasks p(T). More formally, we aim to solve the problem: min, where L i is a loss function for the i-th task T i that we want to minimize. To obtain a model that solves a specific task from the task distribution p(T), we define a task-conditioned model f θ (y|x, z i), with input x, output y, and encoding z i for task T i, which could be provided as a one-hot vector or in any other form. 3 MULTI-TASK LEARNING VIA GRADIENT SURGERY While the multi-task problem can in principle be solved by simply applying a standard single-task algorithm with a suitable task identifier provided to the model or a simple multi-head or multi-output model, a number of prior works (; a;) have found this learning problem to be difficult, especially in the reinforcement learning setting. We hypothesize that one of the main challenges of multi-task learning can be characterized by conflicting and thrashing gradients, and find that this can significantly impede learning progress, especially when combined with iterative data collection. We identify possible causes for this problem and propose a simple and general approach to mitigate it. We hypothesize that a key optimization issue in multi-task learning arises when gradients from multiple tasks are in conflict with one another, i.e. when gradients point away from one another. More specifically, we hypothesize that such conflict may lead to gradient thrashing. Concretely, gradient Figure 2: Visual depiction of conflicting gradients and PCGrad. In (a), we see that tasks A and B have conflicting gradient directions, which can lead to destructive interference and unstable learning. In (b), we illustrate the PCGrad algorithm in cases where gradients are conflicting. PCGrad projects the gradient of task A onto the normal vector of task B's gradient. In (c), we show that tasks with non-conflicting gradients are not altered under PCGrad, thereby keeping tasks with constructive interference. thrashing refers to the phenomenon where a large gradient for one task changes the parameter vectors in a way that substantially decreases performance on another task. Since worse performance typically leads to larger gradients, this in alternating gradient directions, where, at the next iteration, the second task will have large gradients that dominate and reduce performance on the former task. This issue can be particularly pronounced for neural network optimization, since neural network loss landscapes are known to resemble long narrow valleys , where the gradient perpendicular to the direction of the valley will be small. We aim to study this hypothesis through two toy examples. First, consider the two-dimensional optimization landscape illustrated in Fig. 1a, where the landscape for each task objective corresponds to a deep and curved valley (Fig. 1b and 1c). The optima of this multi-task objective correspond to where the two valleys meet. More details on the optimization landscape are in Appendix B. We observe that the gradient thrashing hypothesis is consistent with what we observe when running Adam on this landscape in Fig. 1d, where we observe that Adam does not traverse one valley towards the other, preventing it from reaching an optimum. We also aim to detect if a similar phenomenon occurs in multi-task learning with a neural network with thousands of parameters on a toy regression problem. To measure the extent of gradient thrashing, we plot the cosine similarity between the gradients of two tasks throughout the beginning of learning in Fig. 4 (left). We indeed observe a significant level of gradient thrashing at every iteration, where the cosine similarity varies between −0.75 and 0.75 at a very high frequency. Motivated by these observations, we develop an algorithm that aims to alleviate the optimization challenges caused by gradient thrashing by preventing such gradient conflict between tasks. We aim to prevent gradient thrashing by directly altering the gradients themselves, i.e. through "gradient surgery." To be maximally effective and maximally applicable, we must perform surgery in a way that still allows for positive interactions between the task gradients and does not introduce any assumptions on the form of the model. We start by first detecting whether two gradients are in conflict, by measuring whether they point away from one another. More concretely, we characterize two tasks as conflicting for the current parameter setting if they yield a negative cosine similarity between their respective gradients. The goal of PCGrad is to modify the gradients for each task so as to minimize negative conflict with other task gradients, which will in turn mitigate gradient thrashing. To deconflict gradients during optimization, PCGrad adopts a simple procedure: if the gradients between two tasks are in conflict, i.e. their cosine similarity is negative, we project the gradient from one task onto the normal plane of the gradient of the other task. This amounts to removing the conflicting component of the gradient for the task, thereby reducing the amount of destructive gradient interference between tasks. A pictorial description of this idea is shown in Fig. 2. Suppose the gradient for task T i is g i, and the gradient for task T j is g j. PCGrad proceeds as follows: First, it determines whether g i conflicts with g j by computing the cosine similarity between vectors g i and g j, where negative values indicate conflicting gradients. If the cosine similarity is negative, we replace g i by its projection onto the normal plane of g j: g i = g i − gi·gj gj 2 g j. If the gradients are not in conflict, i.e. cosine similarity is non-negative, the original gradient g i remains unaltered. PCGrad repeats this process across all of the other tasks sampled in random order from the current Compute cosine similarity between gi as gj as cos(φij) = if cos(φij) < 0 then 8: // Subtract the projection of gi onto gj 9: end if 10: end for 11: Store g proj i = gi 12: end for 13: return update ∆θ = i g proj i batch T j ∀ j = i, ing in the gradient g proj i that is applied for task T i. We perform the same procedure for all tasks in the batch to obtain their respective gradients. The full update procedure is described in Algorithm 1 and a discussion on using a random task order is included in Appendix D. This procedure, while simple to implement, ensures that the gradients that we apply for each task per batch interfere minimally with the other tasks in the batch, mitigating the thrashing gradient problem, producing a variant on standard first-order gradient descent in the multi-objective setting. In practice, the PCGrad gradient surgery method can be combined with any gradient-based optimizer, including commonly used methods such as SGD with momentum and Adam , by simply passing the computed update to the respective optimizer instead of the original gradient. Our experimental verify the hypothesis that this procedure reduces the problem of thrashing gradients, and find that, as a , learning progress is substantially improved. Finally, we analyze the convergence of this procedure in Theorem 1 in the two-task setting, to ensure that the procedure is sensible under the standard assumptions in optimization. Theorem 1. Consider two task loss functions L is a multi-task objective. Let φ be the angle between ∇L 1 (θ) and ∇L 2 (θ). Suppose L is differentiable and that its gradient is Lipschitz continuous with constant L > 0, i.e. we have ||∇L(θ 1) − ∇L(θ 2)|| 2 ≤ L||θ 1 − θ 2 || 2 for any θ 1, θ 2. Then, the PCGrad update rule with step size t ≤ 1 L will converge to either a location in the optimization landscape where cos(φ) = −1 or the optimal value L(θ *). Proof. See Appendix A. Theorem 1 states that application of the PCGrad update in the two-task setting with a convex and Lipschitz multi-task loss function L leads to convergence to either the minimizer of L or a potentially sub-optimal objective value. A sub-optimal solution occurs when the cosine similarity between the gradients of the two tasks is −1, i.e. the gradients directly conflict, leading to zero gradient after applying PCGrad. However, in practice, since we are using SGD, which is a noisy estimate of the true batch gradients, the cosine similarity between the gradients of two tasks in a minibatch is unlikely to be −1, thus avoiding this scenario. We apply PCGrad to both supervised learning and reinforcement learning problem settings with multiple tasks or goals. In this section, we discuss the practical instantiations of PCGrad in those settings. Further implementation details are included in Section 6. In multi-task supervised learning, each task. The objective for each task in this supervised setting is then defined as, where z i is a one-hot encoding of task T i. At each training step, we randomly sample a batch of data points B from the whole dataset i D i and then group the sampled data with the same task encoding into small batches denoted as B i for each T i represented in B. We denote the set of tasks appearing in B as B T. After sampling, we precompute the gradient of each task in B T as Given the set of precomputed gradients ∇ θ L i (f θ), we also precompute the cosine similarity between all pairs of the gradients in the set. Using the pre-computed gradients and their similarities, we can obtain the PCGrad update by following Algorithm 1, without re-computing task gradients nor backpropagating into the network. Since the PCGrad procedure is only modifying the gradients of shared parameters in the optimization step, it is model-agnostic and can be readily applied to any architecture designed for supervised multi-task learning. In Section 6, we combine PCGrad with two state-of-the-art architectures for multi-task learning, which leads to noticeable improvement over their original performance. For multi-task reinforcement learning, PCGrad can be readily applied to policy gradient methods by directly updating the computed policy gradient of each task, following Algorithm 1, analogous to the supervised learning setting. For actor-critic algorithms, it is also straightforward to apply PCGrad: we simply replace the task gradients for both the actor and the critic by their gradients computed via PCGrad. Hence, PCGrad can be readily incorporated into a variety of model-free RL algorithms. When applying PCGrad to goal-conditioned RL, we represent p(T) as a distribution of goals and let z i be the encoding of a goal. Similar to the multi-task supervised learning setting discussed above, PCGrad may be combined with various architectures designed for multi-task and goal-conditioned RL , where PCGrad operates on the gradients of shared parameters, leaving task-specific parameters untouched. In our experiments, we apply PCGrad to the soft actor-critic (SAC) algorithm , a recently proposed off-policy actor-critic algorithm that has shown significant gains in sample efficiency and asymptotic performance across many different domains. In SAC, we employ a Qlearning style gradient to compute the gradient of the Q-function network, Q φ (s, a, z i), often known as the critic, and a reparameterization-style gradient to compute the gradient of the policy network π θ (a|s, z i), often known as the actor. For sampling, we instantiate a set of replay buffers {D i} Ti∼p(T). Training and data collection are alternated throughout training. During a data collection step, we run the policy π θ on all the tasks T i ∼ p(T) to collect an equal number of paths for each task and store the paths of each task T i into the corresponding replay buffer D i. At each training step, we sample an equal amount of data from each replay buffer D i to form a stratified batch. For each task T i ∼ p(T), the parameters of the critic θ are optimized to minimize the soft Bellman residual: where γ is the discount factor,φ are the delayed parameters, and α is a learnable temperature that automatically adjusts the weight of the entropy term. For each task T i ∼ p(T), the parameters of the policy π θ are trained to minimize the following objective We compute and apply PCGrad to both following Algorithm 1. In the context of SAC specifically, we further study how the temperature α should be adjusted. If we use a single learnable temperature for adjusting entropy of the multi-task policy π θ (a|s, z i), SAC may stop exploring once all easier tasks are solved, leading to poor performance on tasks that are harder or require more exploration. To address this issue, we propose to learn the temperature on a per-task basis, i.e. using a parametrized model to represent α ψ (z i) (which we abbreviate as PA for per-task alpha). This allows the method to control the entropy of π θ (a|s, z i) per-task. We optimize the parameters of α ψ (z i) using the same constrained optimization framework as in. Algorithms for multi-task learning typically consider how to train a single model that can solve a variety of different tasks (; ;). The multi-task formulation has been applied to many different settings, including supervised learning (; ; ; ;) and reinforcement-learning , as well as many different domains, such as vision (; a; ; ;), language (; ; ;) and robotics; ). While multi-task learning has the promise of accelerating acquisition of large task repertoires, in practice it presents a challenging optimization problem, which has been tackled in several ways in prior work. A number of architectural solutions have been proposed to the multi-task learning problem based on multiple modules or paths (; ; b; b; ; ;), or using attention-based architectures . Our work is agnostic to the model architecture and can be combined with prior architectural approaches in a complementary fashion. A different set of multi-task learning approaches aim to decompose the problem into multiple local problems, often corresponding to each task, that are significantly easier to learn, akin to divide and conquer algorithms a;;;;. Eventually, the local models are combined into a single, multi-task policy using different distillation techniques (outlined in . In contrast to these methods, we propose a simple and cogent scheme for multi-task learning that allows us to learn the tasks simultaneously using a single, shared model without the need for network distillation. Similarly to our work, a number of prior approaches have observed the difficulty of optimization in the multi-task learning setting (; ; b;). Our work, in contrast to many of these optimization schemes, suggests that the challenge in multi-task learning may be attributed to the problem of gradient thrashing, which we address directly by introducing a simple and practical algorithm that de-conflicts gradients from different tasks. Prior work alternatively proposes a gradient-based multi-objective optimization problem for multi-task learning to address the problem of optimizing possibly conflicting objectives. As noted in Alg 2 in , it learns a constant scaling factor for per-task gradient to avoid conflicting, while our method corrects both the scaling factor and the direction of per-task gradient, which can more effectively deconflict gradients. Prior work has also used the cosine similarity between gradients to define when an auxiliary task might be useful for single-task learning . We similarly use cosine similarity between gradients to determine if the gradients between a pair of tasks are in conflict. , we use this measure of gradient conflict as a part of gradient surgery in the context of multi-task learning applications. A number of works in continual learning have studied how to make gradient updates that do not adversely affect other tasks by projecting the gradients into a space that do not conflict with previous tasks . Those methods focus on the continual learning setting, and either need to solve for the gradient projections using quadratic programming , or only projecting the gradient onto the normal plane of the average of the gradients of past tasks . In contrast, our work focuses on multi-task learning, does not require solving any QP, and iteratively projects the gradients of each task onto the normal plane of the gradients of each of the other tasks instead of averaging. Finally, our method is distinct from and solves a different problem than the projected gradient method (Calamai & Moré, 1987), which is an approach for constrained optimization that projects gradients onto the constraint manifold. The goal of our experiments is to study the following questions: Are conflicting gradients a major factor in making optimization for multi-task learning challenging? Does PCGrad make the optimization problems easier for various multi-task learning problems including supervised, reinforcement, and goal-conditioned reinforcement learning settings across different task families? Can PCGrad be combined with other multi-task learning approaches to further improve performance? 6.1 EXPERIMENTAL SETUP To evaluate our method experimentally, we consider both a multi-task supervised learning and a multi-task reinforcement learning problem setup. For supervised learning, we first consider the MultiMNIST dataset , which contains two tasks: classifying the digit on the top left and on the bottom right in an overlaid image. Beyond digit classification, we also use the CIFAR-100 dataset where each of the 20 label superclasses are treated as distinct tasks, following. We also conduct experiments on the NYUv2 dataset , which consists of RGB-D indoor scene images. , we evaluate our method on 3 tasks: 13-class semantic segmentation, depth estimation, and surface normal prediction. In the case of multi-task reinforcement learning, we evaluate our algorithm on the recently proposed Meta-World benchmark . This benchmark includes a variety of simulated robotic manipulation tasks contained in a shared, table-top environment with a simulated Sawyer arm (visualized as the "Push" environment in Fig. 3). In particular, we use the multi-task benchmarks MT10 and MT50, which consists of the 10 tasks and 50 tasks respectively depicted in Fig. 3 that require diverse strategies to solve them, which makes them difficult to optimize jointly with a single policy. Note that MT10 is a subset of MT50. To evaluate goal-conditioned RL scenarios, we consider goal-conditioned robotic pushing with a Sawyer robot. This domain is representative of challenges in learning goal-conditioned policies over a wide distribution of goals. For details on the experimental set-up and model architectures see Appendix E. To answer question, we consider a simple regression problem, where each task is regressing the input to the output of a sine function. The amplitude and the phase of each task are varied. We construct 10 tasks with the amplitude uniformly sampled from the range and is concatenated with the one-hot task encoding. For training, we use a 3-layer fully-connected neural network with 100 hidden units. We compare the performance of the network trained with Adam and the network trained with Adam with PCGrad-modified gradients while plotting the cosine similarity between a pair of tasks during training as shown in Figure 4. The plot on the left in Figure 4 demonstrates that the cosine similarity of Adam gradients between a pair of tasks has high variance, which leads to the gradient thrashing problem, while the cosine similarity of the gradient projected by PCGrad yields positive values diminishing the conflicting-gradients problem. As shown in the plot on the right in Figure 4, Adam with PCGrad leads to faster learning over Adam, which implies that gradient thrashing is indeed a problem in multi-task optimization and reducing it can in considerable performance boost. To answer question, we perform experiments on three standard multi-task supervised learning datasets: MultiMNIST, multi-task CIFAR-100 and NYUv2. We include the on MultiMNIST in Appendix C. For CIFAR-100, we follow to treat 20 coarse labels in the dataset as distinct tasks and create a dataset with 20 tasks and 2500 training instances as well as 500 test instances per task. We combine PCGrad with a powerful multi-task learning architecture, routing networks (; 2019), by simply projecting gradients of the shared parameters in routing networks. As shown in Table 1, applying PCGrad to a single network achieves 71% classification accuracy, which outperforms most of the prior methods such as independent training and cross-stitch (b). Though routing networks achieve better performance than PCGrad on its own, PCGrad is complementary to routing networks and combining PCGrad with routing networks leads to a 2.8% absolute improvement in test accuracy averaged over 3 runs. We also combine PCGrad with another state-of-art multi-task learning algorithm, MTAN , and evaluate the performance on a more challenging indoor scene dataset, NYUv2, which contains 3 tasks as described in Section 6.1. We compare MTAN with PCGrad to a list of methods mentioned in Section 6.1, where each method is trained with three different weighting schemes as in , equal weighting, weight uncertainty (a), and DWA . We only run MTAN with PCGrad with weight uncertainty as we find weight uncertainty as the most effective scheme for training MTAN. The comparing Cross-Stitch, MTAN and MTAN + PCGrad are presented in Table 2 while the full comparison can be found in Table 4 in the Appendix E.3. MTAN with PCGrad is able to achieve the best scores in 8 out of the 9 categories where there are 3 categories per task. Our multi-task supervised learning demonstrate that PCGrad can be seamlessly combined with state-of-art multi-task learning architectures and further improve their on established supervised multi-task learning benchmarks. To answer question, we test all methods on 10 and 50 manipulation tasks respectively shown in Figure 3. At each data collection step, we collect 600 samples for each task, and at each training step, % accuracy task specific-1-fc 42 task specific-all-fc 49 cross stitch-all-fc (b) 53 routing-all-fc + WPL 74.7 independent 67.7 PCGrad (ours) 71 routing-all-fc + WPL + PCGrad (ours) 77.5 (b). Performance of other methods as reported in . we sample 128 datapoints per task from corresponding replay buffers. The are shown in the two plots on the left in Figure 5. We measure success according to the metrics used in the Meta-World benchmark where the reported the success rates are averaged across tasks. For all methods, we apply PA as discussed in Section 4 to learn a separate alpha term per task as the task encoding in MT10 and MT50 is just a one-hot encoding. PCGrad combined with SAC learns all tasks with the best data efficiency and successfully solves all of the 10 tasks in MT10 and about 70% of the 50 tasks in MT50. Training a single SAC policy and a multi-head policy turns out to be unable to acquire half of the skills in both MT10 and MT50, suggesting that eliminating gradient interference across tasks can significantly boost performance of multi-task RL. Training independent SAC agents is able to eventually solve all tasks in MT10 and 70% of the tasks in MT50, but requires about 2 millions and 15 millions more samples than PCGrad with SAC in MT10 and MT50 respectively, implying that applying PCGrad can in leveraging shared structure among tasks that expedites multi-task learning. As noted by , these tasks involve fairly distinct behavior motions, which makes learning all of them with a single policy challenging as demonstrated by poor baseline performance. The ability to learn these tasks together opens the door for a number of interesting extensions to meta-learning, goal conditioned RL and generalization to novel task families. We present the of PCGrad on goal-conditioned RL in the following subsection. We also provide an ablation study on the importance of correcting the gradient direction and scaling the gradient magnitudes in PCGrad. We construct two variants of PCGrad: only applying the gradient direction corrected with PCGrad while keeping the gradient magnitude unchanged and only applying the gradient magnitude computed by PCGrad while keeping the gradient direction unchanged. As shown in the plot on the left in Figure 6, both variants perform worse than PCGrad and the variant where we only vary the gradient magnitudes is much worse than PCGrad. We also compare PCGrad to a prior method GradNorm , which scales the magnitude of gradients of all the tasks. As shown in the plot on the right in Figure 6, PCGrad significantly outperforms GradNorm. We also notice that the variant of PCGrad where only the gradient magnitudes change gets comparable to GradNorm, which suggests that its important to modify both the gradient directions and magnitudes to eliminate interference and achieve good multi-task learning . Figure 6: Ablation study on only using the magnitude and the direction of the gradients modified by PCGrad (left) and comparison between PCGrad and GradNorm (right). PCGrad outperforms both ablations and GradNorm with a large margin, indicating the importance of modifying both the gradient directions and magnitudes in multi-task learning. For our goal-conditioned RL evaluation, we use the robot-pushing environment described in Sec. 6.1 where the goals are represented as the concatenations of the initial positions of the puck to be pushed and the its goal location, both of which are uniformly sampled (details in Appendix E.2). We also apply PA as discussed in Section 4 to predict the temperature for entropy term given the goal. We summarize the in the plot on the right in Figure 5. PCGrad with SAC and PA achieves the best performance in terms of average distance to the goal position, while PCGrad with SAC improves over the baseline and a vanilla SAC agent is struggling to successfully accomplish the task. This suggests that PCGrad is able to ease the RL optimization problem also when the task distribution is continuous. In this work, we identified one of the major challenges in multi-task optimization: conflicting gradients across tasks. We proposed a simple algorithm (PCGrad) to mitigate the challenge of conflicting gradients via "gradient surgery". PCGrad provides a simple way to project gradients to be orthogonal in a multi-task setting, which substantially improves optimization performance, since the task gradients are prevented from negating each other. We provide some simple didactic examples and analysis of how this procedure works in simple settings, and subsequently show significant improvement in optimization for a variety of multi-task supervised learning and reinforcement learning problems. We show that, once some of the optimization challenges of multi-task learning are alleviated by PCGrad, we can obtain the hypothesized benefits in efficiency and asymptotic performance that are believed to be possible in multi-task settings. While we studied multi-task supervised learning and multi-task reinforcement learning in this work, we suspect the problem of conflicting gradients to be prevalent in a range of other settings and applications, such as meta-learning, continual learning, multi-goal imitation learning , and multi-task problems in natural language processing applications . Due to its simplicity and model-agnostic nature, we expect that applying PCGrad in these domains to be a promising avenue for future investigation. Further, the general idea of gradient surgery may be an important ingredient for alleviating a broader class of optimization challenges in deep learning, such as the challenges in the stability challenges in two-player games and multi-agent optimizations . We believe this work to be a step towards simple yet general techniques for addressing some of these challenges. Proof. We will use the shorthand || · || to denote the L 2 -norm and ∇L = ∇ θ L, where θ is the parameter vector. Let g 1 = ∇L 1, g 2 = ∇L 2, and φ be the angle between g 1 and g 2. At each PCGrad update, we have two cases: cos(φ) ≥ 0 or cos(φ < 0). If cos(φ) ≥ 0, then we apply the standard gradient descent update using t ≤ 1 L, which leads to a strict decrease in the objective function value L(φ) unless ∇L(φ) = 0, which occurs only when θ = θ * . In the case that cos(φ) < 0, we proceed as follows: Our assumption that ∇L is Lipschitz continuous with constant L implies that ∇ 2 L(θ) − LI is a negative semidefinite matrix. Using this fact, we can perform a quadratic expansion of L around L(θ) and obtain the following inequality: Now, we can plug in the PCGrad update by letting θ We then get: (Expanding, using the identity (Expanding further and re-arranging terms) (Note that cos(φ) < 0 so the final term is non-negative) Plugging this into the last expression above, we can conclude the following: 2 will always be positive unless ∇L(θ) = 0. This inequality implies that the objective function value strictly decreases with each iteration where cos(φ) > −1. Hence repeatedly applying PCGrad process can either reach the optimal value L(θ) = L(θ *) or cos(φ) = −1, in which case Note that this only holds when we choose t to be small enough, i.e. t ≤ 1 L. To produce the 2D optimization visualizations in Figure 1, we used a parameter vector θ = [θ 1, θ 2] ∈ R 2 and the following task loss functions:. We initialized θ = [0.5, −3] and performed 500,000 gradient updates to minimize L using the Adam optimizer with learning rate 0.001. We compared using Adam for each update to using Adam in conjunction with the PCGrad method presented in Section 3.2. Following the same set-up in , for each image, we sample a different one uniformly at random. Then we put one of the image on the top left and the other one on the bottom right. The two tasks in the multi-task learning problem are to classify the digits on the top left (task-L) and bottom right (task-R) respectively. We construct such 60K examples. We combine PCGrad with the same backbone architecture used in and compare its performance to by running the open-sourced code provided in . As shown in Table 3, our method 0.13% and 0.55% improvement over As stated on line 4 in Algorithm 1, we sample the tasks from the batch and randomly shuffle the order of the tasks before performing the update steps in PCGrad. With random shuffling, we make PCGrad symmetric w.r.t. the task order in expectation. In Figure 7, we observe that PCGrad with a random task order achieves better performance between PCGrad with a fixed task order in the setting of MT50 where the number of tasks is large and the conflicting gradient phenomenon is much more likely to happen. Figure 7: Ablation study on using a fixed task order during PCGrad. PCGrad with a random task order does significantly better PCGrad with a fixed task order in MT50 benchmark. For our CIFAR-100 multi-task experiment, we adopt the architecture used in , which is a convolutional neural network that consists of 3 convolutional layers with 160 3 × 3 filters each layer and 2 fully connected layers with 320 hidden units. As for experiments on the NYUv2 dataset, we follow to use SegNet as the backbone architecture. Our reinforcement learning experiments all use the SAC algorithm as the base algorithm, where the actor and the critic are represented as 6-layer fully-connected feedforward neural networks for all methods. The numbers of hidden units of each layer of the neural networks are 160, 300 and 200 for MT10, MT50 and goal-conditioned RL respectively. We use five algorithms as baselines in the CIFAR-100 multi-task experiment: task specific-1-fc : a convolutional neural network shared across tasks except that each task has a separate last fully-connected layer, task specific-1-fc : all the convolutional layers shared across tasks with separate fully-connected layers for each task, cross stitch-all-fc (b): one convolutional neural network per task along with cross-stitch units to share features across tasks, routing-all-fc + WPL : a network that employs a trainable router trained with multi-agent RL algorithm (WPL) to select trainable functions for each task, independent: training separate neural networks for each task. For comparisons on the NYUv2 dataset, we consider 5 baselines: Single Task, One Task: the vanilla SegNet used for single-task training, Single Task, STAN : the single-task version of MTAN as mentioned below, Multi-Task, Split, Wide / Deep : the standard SegNet shared for all three tasks except that each task has a separate last layer for final task-specific prediction with two variants Wide and Deep specified in , Multi-Task Dense: a shared network followed by separate task-specific networks, Multi-Task Cross-Stitch (b): similar to the baseline used in CIFAR-100 experiment but with SegNet as the backbone, MTAN : a shared network with a soft-attention module for each task. On the multi-task and goal-conditioned RL domain, we apply PCGrad to the vanilla SAC algorithm with task encoding as part of the input to the actor and the critic as described in Section 4 and compare our method to the vanilla SAC without PCGrad and training actors and critics for each task individually (Independent). We use the pushing environment from the Meta-World benchmark as shown in Figure 3. In this environment, the Table 4: We present the full on three tasks on the NYUv2 dataset: 13-class semantic segmentation, depth estimation, and surface normal prediction . #P shows the total number of network parameters. We highlight the best performing combination of multi-task architecture and weighting in bold. The top validation scores for each task are annotated with boxes. The symbols indicate prior methods: *: (a), †: , ‡: (b). Performance of other methods taken from .

Represent this text to find a much much longer article about it with say hundreds-thousands of words

We develop a simple and general approach for avoiding interference between gradients from different tasks, which improves the performance of multi-task learning in both the supervised and reinforcement learning domains.

Filter level sparsity emerges implicitly in CNNs trained with adaptive gradient descent approaches due to various phenomena, and the extent of sparsity can be inadvertently affected by different seemingly unrelated hyperparameters.