Nuclear Norm Regularization for Deep Learning

Authors: christopher  scarvelis
justin  solomon
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-09 23:47:49

Abstract: Penalizing the nuclear norm of a function's Jacobian encourages it to locally behave like a low-rank linear map. Such functions vary locally along only a handful of directions, making the Jacobian nuclear norm a natural regularizer for machine learning problems. However, this regularizer is intractable for high-dimensional problems, as it requires computing a large Jacobian matrix and taking its SVD. We show how to efficiently penalize the Jacobian nuclear norm using techniques tailor-made for deep learning. We prove that for functions parametrized as compositions $f = g \circ h$, one may equivalently penalize the average squared Frobenius norm of $Jg$ and $Jh$. We then propose a denoising-style approximation that avoids the Jacobian computations altogether. Our method is simple, efficient, and accurate, enabling Jacobian nuclear norm regularization to scale to high-dimensional deep learning problems. We complement our theory with an empirical study of our regularizer's performance and investigate applications to denoising and representation learning.

Main paper: 
 introduction: 
 Building models that adapt to the structure of their data is a key challenge in machine learning. As real-world data typically concentrates on low-dimensional manifolds, a good model f should only be sensitive to changes to its inputs along the data manifold. One may encourage this behavior by regularizing f so that its Jacobian Jf [x] has low rank. This causes f to locally behave like a low-rank linear map and therefore be locally constant in directions that are in the kernel of Jf [x] .How should we regularize f so that its Jacobians have low rank? Directly penalizing rank (Jf [x] ) during training is challenging, as the rank function is not differentiable. In light of this, the nuclear norm ∥Jf [x]∥ * is an appealing alternative: Being the L 1 norm of a matrix's singular values, it is the tightest convex relaxation of the rank function, and as it is differentiable almost everywhere, it can be included in standard deep learning pipelines.One critical flaw in this strategy is its computational cost. The nuclear norm of a matrix is the sum of its singular values, so to penalize ∥Jf [x]∥ * in training, one must (1) compute the Jacobian of a typically high-dimensional map f : R n → R m , (2) take the singular value decomposition (SVD) of this m × n matrix, (3) sum its singular values, and (4) differentiate through each of these operations. The combined cost of these operations is prohibitive for high-dimensional data. Consequently, nuclear norm regularization has yet to be widely adopted by the deep learning community.This work shows how to efficiently penalize the Jacobian nuclear norm ∥Jf [x]∥ * using techniques tailor-made for deep learning. We first show that parametrizing f as a composition f = g • h -a feature common to all deep learning pipelines -allows one to replace the expensive nuclear norm ∥Jf [x]∥ * with two squared Frobenius norms ∥Jh[x]∥ 2 F and ∥Jg[h(x)]∥ 2 F , which admit an elementwise computation that avoids a costly SVD. We prove that the resulting problem is exactly equivalent to the original problem with the nuclear norm penalty. We in turn approximate this by a denoising-style objective that avoids the Jacobian computation altogether. Our method is simple, efficient, and accurate -both in theory and in practice -and enables Jacobian nuclear norm regularization to scale to high-dimensional deep learning problems.We complement our theoretical results with an empirical study of our regularizer's performance on synthetic data. As the Jacobian nuclear norm has seldom been used as a regularizer in deep learning, we propose applications of our method to unsupervised denoising, where one trains a denoiser given a dataset of noisy images without access to their clean counterparts, and to representation learning. Our work makes the Jacobian nuclear norm a feasible component of deep learning pipelines, enabling users to learn locally low-rank functions unencumbered by the heavy cost of naïve Jacobian nuclear norm regularization.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: Our main claim is to have proven Theorem 3.1 -which we do in Appendix A.2. The simplicity of our method is apparent -one can implement our regularizer in a few lines of code -and its efficiency follows directly from eliminating the need to compute SVDs or Jacobian matrices. We demonstrate our method's accuracy in Section 4. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 Nuclear norm regularization. As penalizing the nuclear norm in matrix learning problems encourages low-rank solutions, nuclear norm regularization (NNR) has been widely used throughout machine learning. Rennie and Srebro [2005] propose NNR for collaborative filtering, where one attempts to predict user interests by aggregating incomplete information from a large pool of users. Candès et al. [2011] introduce robust PCA, which decomposes a noisy matrix into low-rank and sparse parts and uses nuclear norm regularization to learn the low-rank part. Cabral et al. [2013] , Dai et al. [2014] use NNR to regularize the ill-posed structure-from-motion problem, which recovers a 3D scene from a set of 2D images.A line of work beginning with Candès and Recht [2012] takes inspiration from compressed sensing and studies the conditions under which a low-rank matrix can be perfectly recovered from a small sample of its entries via nuclear norm minimization. This work was followed by Candès and Tao [2010] , Keshavan et al. [2009] , Recht [2011] , which progressively sharpen the bounds on the number of samples required for exact recovery. In parallel, Cai et al. [2010] propose singular value thresholding (SVT), a simple algorithm for approximate nuclear norm minimization that avoids solving a costly semidefinite program as with earlier algorithms. However, SVT still requires computing a singular value decomposition in each iteration, which is an onerous requirement for large matrices. Motivated by this challenge, Rennie and Srebro [2005] show how to convert a nuclear norm-regularized matrix learning problem into an equivalent non-convex problem involving only squared Frobenius norms. Our work generalizes their method to non-linear learning problems.Jacobian regularization in deep learning. Sokolić et al. [2017] , Varga et al. [2017] , Hoffman et al. [2020] penalize the spectral and Frobenius norms of a neural net's Jacobian with respect to its inputs to improve classifier generalization, particularly in the low-data regime. Similarly, Jakubovitz and Giryes [2018] fine-tune neural classifiers with Jacobian Frobenius norm regularization to improve adversarial robustness. Unlike our work, these papers do not consider nuclear norm regularization.Neural ODEs [Chen et al., 2018] parametrize functions as solutions to initial value problems with neural velocity fields. Ensuring that the learned dynamics are well-conditioned minimizes the number of steps required to accurately solve these ODEs. To this end, Finlay et al. [2020] penalize the squared Frobenius norm of the velocity field's Jacobian and observe a tight relationship between the value of this norm and the step size of an adaptive ODE solver. Kelly et al. [2020] extend this approach by regularizing higher-order derivatives as well.Finally, it has been observed as early as Webb [1994] , Bishop [1995] that training a neural net on noisy data approximately penalizes the squared Frobenius norm of the network Jacobian at training data. Inspired by this observation, Vincent et al. [2008 Vincent et al. [ , 2010] ] propose denoising autoencoders for representation learning, and Rifai et al. [2011] propose directly penalizing the squared Frobenius norm of the encoder Jacobian to encourage robust latent representations. Alain and Bengio [2014] show that an autoencoder trained with a penalty on the squared Frobenius norm of its Jacobian learns the score of the data distribution for small regularization values. Recently, Kadkhodaie et al. [2024] employ a similar analysis of the denoising objective to study the generalization of diffusion models.Denoising via singular value shrinkage. While a full survey of the denoising literature is out of scope (see e.g. Elad et al. [2023] ), we highlight a handful of works that employ singular value shrinkage (SVS) to denoise low-rank data corrupted by isotropic noise given a noisy data matrix. SVS denoises a noisy data matrix Y by applying a shrinkage function ϕ to its singular values σ d . This function shrinks small singular values of Y while leaving larger singular values untouched. Shabalin and Nobel [2010] study the conditions under which it suffices to rescale the noisy data's singular values while preserving its singular vectors. Gavish and Donoho [2014] study the optimal hard thresholding policy, which sets all singular values below a threshold to 0 while preserving the rest. Nadakuditi [2014] considers general shrinkage policies whose error is measured by the squared Frobenius distance between the true and denoised data matrix. Gavish and Donoho [2017] then study optimal shrinkage policies under a larger class of error measures. All of these methods assume that the clean data matrix is low-rank and hence that the data globally lies in a low-dimensional subspace. Under the manifold hypothesis, real-world data such as images typically lie on low-dimensional manifolds and hence locally lie in low-dimensional subspaces. Inspired by this observation, we use our method in Section 5 to learn an image denoiser given a dataset of exclusively noisy images.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: Our method's primary limitation is that the Hutchinson-based estimator of the squared Jacobian Frobenius norm introduced in Section 3.3 introduces error that manifests itself as somewhat diffuse boundaries in the 2D experiment in Section 4. We discuss this limitation in Section 6. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
method: 
 In this section, we introduce our method for Jacobian nuclear norm regularization. We first prove that for functions parametrized as compositions f = g • h (which include all non-trivial neural networks), one may replace the expensive nuclear norm penalty ∥Jf [x]∥ * in a learning problem with the average of two squared Frobenius norms 12 ∥Jh[x]∥ 2 F + ∥Jg[h(x)]∥ 2F and obtain an equivalent learning problem. These squared Frobenius norms admit an elementwise computation that avoids a costly SVD in computing ∥Jf [x]∥ * . As the Jacobian computation is itself costly for large neural nets, we use a first-order Taylor expansion and Hutchinson's trace estimator [Hutchinson, 1989] to estimate the Frobenius norm terms. Computing a loss with our regularizer costs as few as two additional evaluations of g and h per sample, allowing our method to scale to large neural networks with high-dimensional inputs and outputs.Any matrix A ∈ R m×n admits a singular value decomposition (SVD) A = U ΣV ⊤ , where U ∈ R m×m and V ∈ R n×n are orthogonal matrices, and Σ ∈ R m×n is a diagonal matrix storing the singular values σ i ≥ 0 of A. The rank of A is equal to its number of non-zero singular values.The nuclear norm ∥A∥ * of a matrix A ∈ R m×n is the sum of its singular values: ∥A∥ * = i σ i . As σ i ≥ 0 by definition, ∥A∥ * is also the L 1 norm of its vector of singular values. Just as L 1 regularization steers learning problems towards solutions with many zero entries, nuclear norm regularization steers matrix learning problems towards low-rank solutions with many zero singular values. For this reason, nuclear norm regularization has seen widespread use in problems ranging from collaborative filtering [Rennie and Srebro, 2005, Candès and Recht, 2012] to robust PCA [Candès et al., 2011] to structure-from-motion [Cabral et al., 2013 , Dai et al., 2014] .However, ∥A∥ * is costly to compute because it requires a cubic-time SVD. Furthermore, even efficient algorithms for nuclear norm minimization such as Cai et al. [2010] 's singular value thresholding require one SVD per iteration, which is prohibitive for very large matrices A. Motivated by this computational challenge, Rennie and Srebro [2005, Lemma 1 ] state that one can compute ∥A∥ * by solving a non-convex optimization problem:EQUATIONFor completeness, we prove this identity in Appendix A.1. Using this result, Rennie and Srebro [2005] show that the following problems are equivalent:EQUATIONwhere r = min (m, n) and ℓ is a generic differentiable loss function. As 1 2 ∥U ∥ 2 F and its derivative ∇ U 1 2 ∥U ∥ 2 F = U can be computed elementwise without a costly SVD, the RHS objective in ( 2) is amenable to gradient-based optimization.Equation (2) enables the use of simple and efficient gradient-based methods for learning a low-rank linear map by parametrizing the matrix A as a composition A = U V ⊤ of linear maps U and V ⊤ . In deep learning, one is interested in learning non-linear functions that are parametrized by compositions of simpler functions. Such functions f are differentiable almost everywhere, so they are locally well-approximated by linear maps specified by their Jacobians Jf [x].Encouraging the learned function to have a low-rank Jacobian is a natural prior: It corresponds to learning a function that locally behaves like a low-rank linear map. Such functions vary locally along only a handful of directions and are constant in the remaining directions. When the training data is supported on a low-dimensional manifold, these directions correspond to the tangents and normals, respectively, to the data manifold. One may implement this low-rank prior on Jf by solving the following optimization problem:EQUATIONwhere ℓ is a generic differentiable loss function and To address these challenges, we first prove a theorem generalizing (2) to non-linear functions. We then show how to avoid computing Jf [x] altogether using a first-order Taylor expansion and Hutchinson's estimator. Our primary contribution is the following result:D(Ω) is a data distribution supported on Ω ⊆ R n . If fTheorem 3.1 Let D(Ω) be a data distribution supported on a compact set Ω ⊆ R n with measure µ that is absolutely continuous with respect to the Lebesgue measure on Ω. Let ℓ ∈ C 1 (R m × R n ) be a continuously differentiable loss function. Then,EQUATIONOn the left-hand side, we learn a function f : R n → R m given fixed input and output dimensions n, m. On the right-hand side, we learn functions h : R n → R d and g : R d → R m with n, m fixed but optimize over the inner dimension d. We prove this theorem in Appendix A.2 and sketch the proof below. Theorem 3.1 shows that by parametrizing f as a composition of g and h -a feature common to all deep learning pipelines -one may learn a locally low-rank function without computing expensive SVDs during training.Proof sketch. We denote the left-hand side objective by E L (f ) and its inf by (L); we denote the right-hand side objective by E R (g, h) and its inf by (R). We prove that (L) ≤ (R) and (R) ≤ (L).(L) ≤ (R) is the easy direction. The basic observation is that iff = g • h, then Jf [x] = Jg[h(x)]Jh[x]by the chain rule. Equation (1) then implies that∥Jf [x]∥ * = min U,V :U V ⊤ =Jf [x] 1 2 ∥U ∥ 2 F + ∥V ∥ 2 F ≤ 1 2 ∥Jg[h(x)]∥ 2 F + ∥Jh[x]∥ 2 F .(R) ≤ (L) is the hard direction. The proof strategy is as follows:1. We begin with a function f m ∈ C ∞ (Ω) such that E L (f m ) is arbitrarily close to its inf over C ∞ (Ω). We use f m to construct parametric families of affine functions g z m , h z m whose composition is a good local approximation to f m in a neighborhood of z ∈ Ω, both pointwise and in terms of the contributions to E R (g z m , h z m ) and E L (f m ), resp., due to x ∈ Ω near z.2. We then stitch together these local approximations to form a sequence of global approximations g k m , h k m to f m . These functions are piecewise affine and hence not regular enough to lie in C ∞ (Ω) as required by the right-hand side of Equation (4).k m , h k m yields a minimizing sequence of C ∞ (Ω) functions g k m,ϵ , h k m,ϵ such that E R (g k m,ϵ , h k m,ϵ ) approaches the inf of E L .While Theorem 3.1 shows how to regularize a learning problem with the Jacobian nuclear norm without a cubic-time SVD, the Jacobian computation incurs a quadratic time and memory cost, which remains heavy for high-dimensional learning problems. To mitigate this issue, the following section shows how to approximate the ∥Jg[h(x)]∥ 2 F and ∥Jh[x]∥ 2 F terms in (4).When f : R n → R m is a function between high-dimensional spaces, Jf [x] is an m × n matrix that is costly to compute and to store in memory. Previous works employing Jacobian regularization for neural networks have noted this issue and proposed stochastic approximations based on Jacobianvector products (JVP) against random vectors [Varga et al., 2017 , Hoffman et al., 2020] . As JVPs may be costly to compute for large neural nets, we propose an alternative stochastic estimator that requires only evaluations of f and analyze its error:EQUATIONSimilar results appear in the ML literature as early as Webb [1994] . Our proof in Appendix A.3 relies on a first-order Taylor expansion and Hutchinson's trace estimator; a similar proof is given by Alain and Bengio [2014] . In practice, we obtain accurate approximations to ∥Jf [x]∥ 2 F by using a small noise variance σ 2 and rescaling the expectation in (5) by 1 σ 2 to compensate. In Section 4, we also show that a single noise sample ϵ ∼ N (0, σ 2 I) suffices in practice.Using this efficient approximation, we obtain the following regularizer:EQUATIONwhere f = g • h. In practice, one may use a single draw of ϵ ∼ N (0, σ 2 I) per training iteration while maintaining good performance on downstream tasks; see e.g. the results in Section 5. In this case, our regularizer R(x; f ) costs merely two additional function evaluations, enabling it to scale to large neural networks acting on high-dimensional data. In Section 4, we show that parametrizing f θ as a neural net and solving infEQUATIONyields good approximations to the solution to (3) for problems where exact solutions are known. We then propose two applications of Jacobian nuclear norm regularization in Section 5.k m , h k m yields a minimizing sequence of C ∞ (Ω) functions g k m,ϵ , h k m,ϵ such that E R (g k m,ϵ , h k m,ϵ ) approaches the inf of E L .Local approximations to f m . To begin, letf m ∈ C ∞ (Ω) be a function that attains E L (f m ) ≤ inf f ∈C ∞ (Ω) E L (f ) + 1 m . Fix some z ∈ Ω, take the thin SVD of Jf m [z] = U m (z)Σ m (z)V m (z)⊤ , and use it to define two parametric families of affine functions:EQUATIONg z m (y) = U m (z) Σ m (z)y + f m (z) -Jf m [z]z.These functions satisfy two key properties:EQUATIONwhere∥R z m (x)∥ 2 ∈ O(∥x -z∥ 2 2) by Taylor's theorem, andη 2 ∥Jg z m [h z m (x)]∥ 2 F + ∥Jh z m [x]∥ 2 F = η 2 ∥ Σ m (z)∥ 2 F + ∥ Σ m (z)∥ 2 F = η∥Jf m [z]∥ * . (15)Using ( 14), we obtainEQUATIONThe continuity of ℓ ensures that we can make ℓ(f m (x) + R z m (x), x) arbitrarily close to ℓ(f m (x), x) by making ∥x -z∥ 2 sufficiently small. Furthermore,EQUATIONand as f m ∈ C ∞ , Jf m is a continuous function, so we can make ∥Jf m [z] -Jf m [x]∥ * arbitrarily small by making ∥x -z∥ sufficiently small.The compositions of these functions g z m , h z m are good local approximations to f m in a neighborhood of z, both pointwise (by ( 14)) and in terms of the energy contributions arising from x ∈ Ω near z (by ( 16) and ( 17)).Global piecewise affine approximations to f m . We now stitch together these local approximations to form a sequence of global approximations g k m , h k m to f m .For each k, fix a set of pointsZ k = {z i } N (k)i=1 and use this set to partition Ω into Voronoi regions V i . Choose the centroids z i such that max x∈Vi ∥x -z i ∥ 2 ≤ ϵ k , for ϵ k > 0 sufficiently small to ensure that |∥Jf m [z i ]∥ * -∥Jf m [x]∥ * | + |ℓ(f m (x) + R zi m (x), x) -ℓ(f m (x), x)| < 1k for all x ∈ V i and for all regions V i . (The compactness of Ω and the uniform continuity of ∥Jf m ∥ * on Ω and ℓ on its domain ensures that we can always find a finite set of centroids with this property.) Then let g k m , h k m be piecewise affine functions such that g k m (x) := g zi m (x) and h k m (x) := h zi m (x) for all x ∈ int(V i ). For all points x on a Voronoi boundary, define g k m , h k m by averaging over the interiors of the Voronoi cells incident on the boundary. This yields:E x∼D(Ω) ℓ(g k m (h k m (x)), x) + η 2 ∥Jg k m [h k m (x)]∥ 2 F + ∥Jh k m [x]∥ 2 F = E x∼D(Ω)   N (k) i=1 [ℓ(f m (x) + R zi m (x)), x) + η∥Jf m [z i ]∥ * ] • 1 [x ∈ V i ]   ≤ E x∼D(Ω)   N (k) i=1 ℓ(f m (x), x) + η∥Jf m [x]∥ * + 1 k • 1 [x ∈ V i ]   = E x∼D(Ω) ℓ(f m (x), x) + η∥Jf m [x]∥ * + 1 k = E L (f m ) + 1 k ≤ inf f ∈C ∞ (Ω) E L (f ) + 1 m + 1 k and hence E R (g k m , h k m ) ≤ inf f ∈C ∞ (Ω) E L (f ) + 1 m + 1 k .Mollifying the piecewise affine approximations. The functions g k m , h k m constructed in the previous section are piecewise affine and hence not regular enough to lie in C ∞ (Ω) as required by (R). We now mollify these functions to yield a minimizing sequence ofC ∞ (Ω) functions g k m,ϵ , h k m,ϵ such that E R (g k m,ϵ , h k m,ϵ ) approaches inf f ∈C ∞ (Ω) E L (f ).We mollify g k m , h k m by convolving them against the standard mollifiers (infinitely differentiable and compactly supported on B(0, ϵ)) to yield a sequence of C ∞ (Ω) functions g k m,ϵ , h k m,ϵ . We need to show thatE R (g k m,ϵ , h k m,ϵ ) ≤ E R (g k m , h k m ) + ψ(ϵ; m, k)for some error ψ(ϵ; m, k) that vanishes as ϵ → 0 for any m, k. We proceed by individually controlling the terms in E R (g k m,ϵ , h k m,ϵ ). 
validation: 
 In this section, we empirically validate our method on a special case of (3) for which closed-form solutions are known. We consider the following problem:EQUATIONwhere τ : R n → R is the indicator function of the unit ball in R n . As f is a scalar-valued function in this problem, Jf [x] is a vector, and ∥Jf [x]∥ * = ∥∇f (x)∥ 2 . This is an instance of the celebrated Rudin-Osher-Fatemi (ROF) model for image denoising [Rudin et al., 1992] . Meyer [2001, p. 36] shows that the exact solution to (8) given this target function τ (x) is f (x) := (1 -nη)τ (x). This is a rescaled indicator function of the unit ball.We parametrize f as a multilayer perceptron (MLP) f θ and solve (8) along with the problem using our regularizer:EQUATIONwhere f θ = g θ • h θ . We approximate the integral over R n by Monte Carlo integration over a box Ω centered at the origin. We experiment with n = 2 and n = 5 and in each case depict results for a small and a large regularization value η. We track the objective values of problems ( 8) and ( 9) and show that they converge to the same value, as predicted by Theorem 3.1. We also track the absolute error of both problem's solutions across training iterations and plot solutions to each problem at convergence for the 2D case. We give full implementation details in Appendix B.1. Figure 2 confirms the accuracy our our method's solutions, which attain absolute error comparable to the neural solutions to (8).Figure 3 depicts the objective values for Problems ( 8) and ( 9) on log scale across training iterations. As predicted by Theorem 3.1, both converge toward nearly identical objective values; the larger gap in Figure 3 (c) is an artifact of the loss magnitudes being smaller and the plot being on log scale. 
applications: 
 Unsupervised denoising. In this section, we apply our regularizer R(x; f θ ) to unsupervised denoising. We learn a denoiser f θ that maps noisy images x + ϵ to clean images x given a training set of noisy images without their corresponding clean images. We motivate the use of our regularizer via a connection with denoising by singular value shrinkage, and demonstrate that our unsupervised denoiser nearly matches the performance of a denoiser trained on clean-noisy image pairs. Singular value shrinkage. A line of work beginning with Shabalin and Nobel [2010] studies denoising by singular value shrinkage (SVS). These works seek to recover a low-rank data matrix X ∈ R D×N of unknown rank r given only a single matrix Y = X + σ ϵ Z of clean data X corrupted by iid white noise Z. Since the clean data is low-rank, the components of Y corresponding to its small singular values contain mostly noise, so SVS denoises Y by applying a shrinkage function ϕ to its singular values σ d . This function shrinks small singular values of Y while leaving larger singular values untouched. For convenience, we denote the denoised matrix by ϕ(Y ).Gavish and Donoho [2017] show that under certain assumptions on the noise Z and data X, one can derive the optimal shrinker ϕ that asymptotically minimizes ∥X -ϕ(Y )∥ 2 F . In Appendix A.4, we show that this optimal shrinker is also the solution to the following problem:EQUATIONwhere we set η = σ 2 ϵ to be equal to the noise variance. This problem is a special case of the following instance of Problem ( 3)inf f θ :R D →R D E y∼D(Ω) 1 2 ∥f θ (y) -y∥ 2 2 + η∥Jf θ [y]∥ * (11)when D(Ω) is an empirical distribution over N noisy training samples and f θ is restricted to be a linear map. Just as (10) yields an optimal shrinker for denoising low-rank data which globally lies in a low-dimensional subspace, we conjecture that solving (11) yields an effective denoiser for manifold-supported data such as images, which locally lie near low-dimensional subspaces -even when trained on noisy images. We test this conjecture by solving (11) using a neural denoiser f θ . To make this problem tractable, we replace ∥Jf θ [y]∥ * with our regularizer R(y; f θ ), which we compute using a single draw of ϵ per training iteration.Experiments. We train our denoiser by solving (11) with D(Ω) being the empirical distribution over 288k noisy images from the Imagenet training set [Russakovsky et al., 2015] . Consequently, We set η = σ 2 when solving (11). We parametrize the denoiser f θ = g θ • h θ as a Unet [Ronneberger et al., 2015] , letting h θ and g θ be its downsampling and upsampling blocks, resp.PSNR (dB) ↑ Imagenet CBSD68 Method σ = 1 σ = 2 σ = 1 σ = 2 BM3D 21.We benchmark our method against a supervised denoiser trained with the usual MSE loss ∥f θ (x + ϵ) -x∥ 2 2 on clean-noisy pairs, Lehtinen et al. [2018] 's Noise2Noise (N2N) method, which requires access to independent noisy copies of each ground truth image during training, and BM3D, a classical unsupervised denoiser [Dabov et al., 2007] . We implement the supervised and N2N denoisers using the same Unet architecture as our denoiser, and train them on the same dataset with the same hyperparameters. We evaluate the denoisers via their average peak signal-to-noise ratio (PSNR) across the CBSD68 dataset [Martin et al., 2001] and across 100 randomly-drawn noisy images from the Imagenet validation set, randomly cropped to 256 × 256. We provide full details for this experiment in Appendix B.2.We report each denoiser's performance in Table 1 and include 1-sigma error bars computed across the test images. Despite being trained exclusively on highly corrupted images, our denoiser nearly matches the performance of an ordinary supervised denoiser at both noise levels and performs comparably to Noise2Noise, which requires independent noisy copies of each ground truth image during training. We further illustrate the comparison in Figure 4 . All neural methods recover substantially more fine detail than the classical BM3D denoiser, particularly at the larger noise level σ = 2. Notably, our method performs nearly as well as a supervised denoiser, despite being trained exclusively on highly corrupted data.We also demonstrate the sparsity-inducing effect of our regularizer on the singular values of Jf θ in Figure 5 , where we plot the Jacobian singular values of our denoiser and a supervised denoiser at a randomly-drawn validation image corrupted with σ = 2 Gaussian noise. We normalize the singular values so that each Jacobian's largest singular value is 1 and depict the singular values on log scale. As expected, our denoiser's Jacobian singular values decay more rapidly than those of the supervised denoiser at the same point.These results show that our regularizer (6) can be used to construct a tractable non-linear generalization of Gavish and Donoho [2017] 's optimal shrinker that performs nearly as well as a supervised denoiser on image denoising tasks, despite being trained exclusively on highly corrupted images. Representation learning. We now apply our method to unsupervised representation learning. We train a deterministic autoencoder consisting of an encoder f θ and a decoder g ϕ on the CelebA dataset [Liu et al., 2015] and approximately penalize the Jacobian nuclear norm ∥J f θ [x]∥ * of the encoder at training data x ∈ D(Ω) using our regularizer R(x; f θ ). This encourages the encoder to locally behave like a low-rank linear map whose image is low-dimensional. One may interpret this as a deterministic autoencoder with locally low-dimensional latent spaces. We demonstrate that the left-singular vectors of the encoder Jacobian J f θ [x] are semantically meaningful directions of variation about training data in latent space. We provide full experimental details in Appendix B.3.To demonstrate our autoencoder's ability to learn meaningful representations, we select an arbitrary training point x and traverse the latent space of our regularized autoencoder and an unregularized baseline along rays of the form z = f θ (x) + αu d θ (x), where u d θ (x) is the d-th left-singular vector of the encoder Jacobian Jf θ [x]. These left-singular vectors form a basis for the image of Jf θ [x] and approximate a basis for the tangent space of the latent manifold at f θ (x). We depict decoded images along this traversal for our regularized autoencoder in Figure 7 and for the baseline in Figure 6 . The traversals in Figure 7 are semantically meaningful. For instance, a traversal along the first singular vector edits the tint of the decoded image, and a traversal along the second singular vector edits the facial expression of the image's subject. The traversals of the unregularized autoencoder's latent space in Figure 6 edit the colors of the decoded image but are unable to control other attributes.We also follow Higgins et al. [2017] and visualize traversals along latent variables of a β-VAE at the same training point x in Figure 8 . While the β-VAE is able to discover meaningful directions of variation, the decoded images are highly diffuse, as is typical of VAEs. In contrast, our autoencoder's reconstructions retain finer details. We conjecture that our model's improved capacity results from our autoencoder's ability to learn a locally low-dimensional latent space without constraining its global structure. 
conclusion: 
 The Jacobian nuclear norm ∥Jf [x]∥ * is a natural regularizer for learning problems, where it steers solutions towards having low-rank Jacobians. Such functions are locally sensitive to changes in their inputs in only a few directions, which is an especially desirable prior for data that is supported on a low-dimensional manifold. However, computing ∥Jf [x]∥ * naïvely requires both evaluating a Jacobian and taking its SVD; the combined cost of these operations is prohibitive for the highdimensional maps f that often arise in deep learning.Our work resolves this computational challenge by generalizing a surprising result (2) from matrix learning to non-linear learning problems. As they rely on parametrizing the learned function f = g •h as a composition of functions g and h, our methods are tailor-made for deep learning, where such parametrizations are ubiquitous. We anticipate that the deep learning community will discover additional applications of Jacobian nuclear norm regularization to make use of our efficient methods.As an efficient implementation of our regularizer relies on estimating the squared Jacobian Frobenius norm using Hutchinson's trace estimator, some error is inevitable. This error manifests itself in slightly diffuse boundaries in the solutions to the Rudin-Osher-Fatemi problem in Figure 1 . However, we do not find this error problematic for high-dimensional applications such as unsupervised denoising as in Section 5. Future implementations of our method may employ more accurate estimators of the squared Jacobian Frobenius norm for applications where accuracy is of paramount concern. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Our method does not have a high risk of misuse. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We cite the datasets, pre-existing algorithms, and code libraries used for our experiments throughout our work. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A.1 Proof of Equation (1)We draw heavy inspiration from a proof by Fritz that appeared on MathOverflow. We will prove that∥A∥ * = min U,V :U V ⊤ =A 1 2 ∥U ∥ 2 F + ∥V ∥ 2 F .First suppose that U, V are matrices such that U V ⊤ = A. By the matrix Hölder inequality,∥A∥ * = ∥U V ⊤ ∥ * ≤ ∥U ∥ F • ∥V ∥ F .By the arithmetic mean-geometric mean (AM-GM) inequality,∥U ∥ F • ∥V ∥ F = ∥U ∥ 2 F • ∥V ∥ 2 F ≤ 1 2 ∥U ∥ 2 F + ∥V ∥ 2 F .Combining these results, we obtain∥A∥ * = ∥U V ⊤ ∥ * ≤ inf U,V :U V ⊤ =A 1 2 ∥U ∥ 2 F + ∥V ∥ 2 F .To see that the inf is attained at ∥A∥ * , take the compact SVDEQUATION1 2 ∥U ∥ 2 F + ∥V ∥ 2 F = 1 2 ∥U A √ Σ∥ 2 F + ∥V A √ Σ∥ 2 F = 1 2 ∥ √ Σ∥ 2 F + ∥ √ Σ∥ 2 F = 1 2 r i=1 σ i + r i=1 σ i = ∥A∥ *This proves Equation ( 1). An alternative proof using somewhat more complex methods is given in Mazumder et al. [2010] . ■A.2 Proof of Theorem 3.1Let D(Ω) be a data distribution with measure µ supported on a compact set Ω ⊆ R n that is absolutely continuously with respect to the Lebesgue measure λ(Ω), and let ℓ ∈ C 1 (R m ) be a continuously differentiable loss function. Let C ∞ (Ω) denote the set of infinitely differentiable functions on Ω. We will show that:EQUATIONWe denote the left-hand side objective by E L (f ) and its inf by (L); we denote the right-hand side objective by E R (g, h) and its inf by (R). We prove that (L) ≤ (R) and (R) ≤ (L).A.2.1 (L) ≤ (R)This is the easy direction. The basic observation is that if we parametrizef = g • h, then Jf [x] = Jg[h(x)]Jh[x], and Srebro's identity (1) tells us that∥Jf [x]∥ * = min U,V :U V ⊤ =Jf [x] 1 2 ∥U ∥ 2 F + ∥V ∥ 2 F ≤ 1 2 ∥Jg[h(x)]∥ 2 F + ∥Jh[x]∥ 2 FThe rest of the proof is book-keeping.We first verify that C ∞ (Ω) is closed under composition by showing that:EQUATION)The ⊆ inclusion is straightforward: given any f ∈ C ∞ (Ω), just choose h ≡ f and g ≡ Id; these functions are clearly in the correct classes. The ⊇ inclusion follows from the fact that the composition of C ∞ functions is C ∞ by the chain rule.This yields:inf f ∈C ∞ (Ω) E L (f ) = inf h∈C ∞ (Ω) g∈C ∞ (h(Ω)) E L (g • h).We now use Srebro's identity as follows:inf h∈C ∞ (Ω) g∈C ∞ (h(Ω)) E L (g • h) = inf h∈C ∞ (Ω) g∈C ∞ (h(Ω)) E x∼D(Ω) [ℓ(g(h((x)), x) + η∥J(g • h)[x]∥ * ] = inf h∈C ∞ (Ω) g∈C ∞ (h(Ω)) E x∼D(Ω) [ℓ(g(h((x)), x) + η∥Jg[h(x)]Jh[x]∥ * ] ≤ inf h∈C ∞ (Ω) g∈C ∞ (h(Ω)) E x∼D(Ω) ℓ(g(h((x)), x) + η 2 ∥Jg[h(x)]∥ 2 F + ∥Jh[x]∥ 2 F = inf h∈C ∞ (Ω) g∈C ∞ (h(Ω)) E R (g • h). Consequently, inf f ∈C ∞ (Ω) E L (f ) ≤ inf h∈C ∞ (Ω) g∈C ∞ (h(Ω)) E R (g • h)and hence (L) ≤ (R).A.2.2 (R) ≤ (L)This is the hard direction. The proof strategy is as follows:1. We begin with a functionf m ∈ C ∞ (Ω) such that E L (f m ) is arbitrarily close to its inf over C ∞ (Ω).We use f m to construct parametric families of functions g z m , h z m whose composition is a good local approximation to f m in a neighborhood of z ∈ Ω, both pointwise and in terms of the energy contributions due to x ∈ Ω near z. This step relies crucially on our ability to construct optimal solutions to the RHS problem in (1).2. We then stitch together these local approximations to form a sequence of global approximations g k m , h k m to f m . These functions are piecewise affine and hence not regular enough to lie in C ∞ (Ω) as required by the right-hand side of Equation (12).E x∼D(Ω) ℓ(g k m,ϵ (h k m,ϵ (x)), x) .We first show that it suffices to prove that∥g k m,ϵ • h k m,ϵ -g k m • h k m ∥ L 1 (Ω,µ) → 0 for any m, k.(Recall that µ is the measure associated with the data distribution D(Ω).)Note that as ℓ ∈ C 1 (R m ) and the image of the compact set Ω under the piecewise affine functiong k m • h k m is bounded, ℓ is L-Lipschitz on g k m (h k m (Ω)).For any sequence of functions f n converging to f in L 1 , we then have:Ω ℓ(f n (x), x)dµ - Ω ℓ(f (x), x)dµ ≤ Ω |ℓ(f n (x), x) -ℓ(f (x), x)| dµ ≤ Ω L∥f n (x) -f (x)∥ 2 dµ = L Ω ∥f n (x) -f (x)∥ 2 dµ = L∥f n -f ∥ L 1 (Ω,µ) → ϵ→0 0 It therefore suffices to show that ∥g k m,ϵ • h k m,ϵ -g k m • h k m ∥ L 1 (Ω,µ) → 0 for any m, k.To this end, we will use the bound∥g k m,ϵ •h k m,ϵ -g k m •h k m ∥ L 1 (Ω,µ) ≤ ∥g k m,ϵ •h k m,ϵ -g k m,ϵ •h k m ∥ L 1 (Ω,µ) +∥g k m,ϵ •h k m -g k m •h k m ∥ L 1 (Ω,µ)and show that each of the RHS terms goes to 0 as ϵ → 0.We begin by controlling∥g k m,ϵ • h k m,ϵ -g k m,ϵ • h k m ∥ L 1 (Ω,µ). The obvious approach is to use the fact that g k m,ϵ is Lipschitz on the bounded domain h k m (Ω) and that ∥h k m,ϵ -h k m ∥ L 1 (Ω,µ) → 0 by standard properties of mollifiers. However, this naïve approach fails because the Lipschitz constant of g k m,ϵ increases as ϵ → 0: We need to bound it in terms of a quantity independent of ϵ. We will instead use the fact that the un-mollified function g k m is piecewise affine and hence locally Lipschitz on the interior of each affine region.Before proceeding with the remainder of the proof, we make a key observation. Given a Voronoi partition of the compact domain Ω into N (k) cells V i , we constructed g k m , h k m as piecewise affine functions such that g k m (x) := g zi m (x) and h k m (x) := h zi m (x) for all x ∈ int(V i ) -and for all points x on a Voronoi boundary, we defined g k m , h k m by averaging over the interiors of the Voronoi cells incident on the boundary. The affine functions g zi m , h zi m were constructed to have the following properties:EQUATIONwhere ∥R z m (x)∥ 2 ∈ O(∥x -z∥ 2 2 ) by Taylor's theorem, andEQUATION)At any x on the interior of a Voronoi cell, g k m and h k m are equal to some affine functions g zi m , h zi m , so by(19), η 2 ∥Jg k m [h k m (x)]∥ 2 F + ∥Jh k m [x]∥ 2 F = η∥Jf m [z i ]∥ * < +∞. In particular, ∥Jg k m [h k m (x)]∥ 2 F and ∥Jh k m [x]∥ 2F must both be well-defined and finite at this x. This can only happen if h k m (x) lies on the interior of a Voronoi cell, since otherwise Jg k m [h k m (x)] would be undefined. Hence h k m maps the interior of Voronoi cells to the interior of Voronoi cells; the contrapositive is that if h k m (x) lies on a Voronoi boundary, then x also lies on a Voronoi boundary. It follows that the preimage of the Voronoi boundaries under h k m is a set of Lebesgue measure zero. Under the absolute continuity hypothesis, this is also a set of µ-measure zero. We use this fact extensively throughout the remainder of the proof.For any m, k, let S k m ⊆ Ω be the set of Voronoi boundaries from the partition we constructed earlier; this Voronoi partition depends on m, k but not on ϵ. We then begin by stripping the Voronoi boundaries from Ω to obtain Ω(m, k) := Ω \ S k m . This only removes a set of µ-measure 0 from Ω and hence doesn't impact the remaining convergence results. Given ϵ, we then partition Ω(m, k) into two subsets:• Let Ω 0 (m, k, ϵ) ⊆ Ω(m, k) be the set of x ∈ Ω(m, k) such that d(x, S k m ) > ϵ.(We use d(p, S) to denote the distance from the point p from the set S.)• Let Ω 1 (m, k, ϵ) be its complement in Ω(m, k): the set of x ∈ Ω(m, k) such that d(x, S k m ) ≤ ϵ.Ω 0 (m, k, ϵ) will be our good set, and Ω 1 (m, k, ϵ) will be a bad set whose measure converges to 0 asϵ → 0. We decompose ∥g k m,ϵ • h k m,ϵ -g k m,ϵ • h k m ∥ L 1 (Ω,µ)as follows:∥g k m,ϵ • h k m,ϵ -g k m,ϵ • h k m ∥ L 1 (Ω,µ) = Ω0(m,k,ϵ) ∥g k m,ϵ (h k m,ϵ (x)) -g k m,ϵ (h k m (x))∥ 2 dµ + Ω1(m,k,ϵ) ∥g k m,ϵ (h k m,ϵ (x)) -g k m,ϵ (h k m (x))∥ 2 dµWe begin by showing that ∥g k m,ϵ (h k m,ϵ (x)) -g k m,ϵ (h k m (x))∥ 2 = 0 for all x ∈ Ω 0 (m, k, ϵ). To this end, first recall that the standard mollifier supported on B(0, ϵ) is defined as follows:η ϵ (y) = C(ϵ) exp 1 ∥ y ϵ ∥ 2 2 -1 for ∥y∥ 2 ≤ 1 0 for ∥y∥ 2 > 1 (20)where C(ϵ) > 0 is chosen to ensure that η ϵ integrates to 1. Now, if x ∈ Ω 0 (m, k, ϵ), then d(x, S k m ) > ϵ and hence B(x, ϵ) is entirely contained in the Voronoi cell V i containing x. (Note that x cannot lie on a Voronoi boundary, as we have stripped the µ-measure zero set of Voronoi boundaries S k m from Ω before constructing Ω 0 (m, k, ϵ) and Ω 1 (m, k, ϵ).) On this ball B(x, ϵ), the linearity of h k m (y) := Σ m (z i )V m (z i ) ⊤ y and the rotational symmetry of the mollifier η ϵ yields:EQUATIONHence for allx ∈ Ω 0 (m, k, ϵ), h k m,ϵ (x) = h k m (x), and it follows that g k m,ϵ (h k m,ϵ (x)) = g k m,ϵ (h k m (x)) and therefore ∥g k m,ϵ (h k m,ϵ (x)) -g k m,ϵ (h k m (x))∥ 2 = 0. We conclude that Ω0(m,k,ϵ) ∥g k m,ϵ (h k m,ϵ (x)) -g k m,ϵ (h k m (x))∥ 2 dµ = 0.To bound the second term, note that the following inequalities hold µ-almost everywhere on Ω:EQUATIONThe second inequality can be derived using Jensen's inequality by first showing the following for all y ∈ h k m (Ω(m, k)):∥g k m,ϵ (y)∥ 2 = ∥ B(y,ϵ) g k m (z) η ϵ (z)dz :=dηϵ(z) ∥ 2 ≤ Jensen B(y,ϵ) ∥g k m (z)∥ 2 dη ϵ (z) ≤ sup y∈h k m (Ω(m,k)) ∥g k m (y)∥ 2 B(y,ϵ) dη ϵ (z) =1 = sup y∈h k m (Ω(m,k)) ∥g k m (y)∥ 2 , which implies that sup y∈h k m (Ω(m,k)) ∥g k m,ϵ (y)∥ 2 ≤ sup y∈h k m (Ω(m,k)) ∥g k m (y)∥ 2 .Returning to (26), we can bound the integral over the bad setΩ1(m,k,ϵ) ∥g k m,ϵ (h k m,ϵ (x)) - g k m,ϵ (h k m (x))∥ 2 dµ as follows: Ω1(m,k,ϵ) ∥g k m,ϵ (h k m,ϵ (x)) -g k m,ϵ (h k m (x))∥ 2 dµ ≤ M (m, k) • µ(Ω 1 (m, k, ϵ))where µ(Ω 1 (m, k, ϵ)) is the measure of Ω 1 (m, k, ϵ). We will now show that µ(Ω 1 (m, k, ϵ)) → 0, which will implyΩ1(m,k,ϵ) ∥g k m,ϵ (h k m,ϵ (x)) -g k m,ϵ (h k m (x))∥ 2 dµ → ϵ→0 0,and therefore allow us to conclude that∥g k m,ϵ • h k m,ϵ -g k m,ϵ • h k m ∥ 2 L 1 (Ω,µ) → ϵ→0 0.Proving that µ(Ω 1 (m, k, ϵ)) → 0. Recall that:EQUATIONwhere d(x, S k m ) denotes the distance from the point x to the union of Voronoi boundaries S k m . Note that this set is a union of cylinders of radius ϵ centered at the Voronoi boundaries S k m . As S k m has Lebesgue measure 0, the Lebesgue measure of the cylinders B(m, k, ϵ) also goes to 0 as the radius ϵ → 0. The absolute continuity of µ then implies that µ(Ω 1 (m, k, ϵ)) → 0 as well.Using these results, we obtainΩ1(m,k,ϵ) ∥g k m,ϵ (h k m,ϵ (x)) -g k m,ϵ (h k m (x))∥ 2 dµ → ϵ→0 0,and therefore conclude that∥g k m,ϵ • h k m,ϵ -g k m,ϵ • h k m ∥ 2 L 1 (Ω,µ) → ϵ→0 0.This completes the proof that∥g k m,ϵ • h k m,ϵ -g k m,ϵ • h k m ∥ L 1 (Ω,µ) → 0 as ϵ → 0. Controlling the second term ∥g k m,ϵ • h k m -g k m • h k m ∥ L 1 (Ω,µ) is easier. Indeed, g k m,ϵ • h k m → g k m • h k m pointwise a.e. as ϵ → 0 by standard properties of mollifiers. Furthermore, the sequence g k m,ϵ • h k m is dominated almost everywhere by the function identically equal (componentwise) tomax x∈Ω(m,k) g k m (h k m (x)); this function is in L 1 (Ω, µ) because Ω is a compact domain. It follows from the dominated convergence theorem that ∥g k m,ϵ • h k m -g k m • h k m ∥ L 1 (Ω,µ) → 0 as ϵ → 0.We have shown that each of the following RHS terms goes to 0 as ϵ → 0:∥g k m,ϵ •h k m,ϵ -g k m •h k m ∥ L 1 (Ω,µ) ≤ ∥g k m,ϵ •h k m,ϵ -g k m,ϵ •h k m ∥ L 1 (Ω,µ) +∥g k m,ϵ •h k m -g k m •h k m ∥ L 1 (Ω,µ) , which allows us to conclude that ∥g k m,ϵ • h k m,ϵ -g k m • h k m ∥ L 1 (Ω,µ) → 0 and consequently E x∼D(Ω) ℓ(g k m,ϵ (h k m,ϵ (x)), x) -E x∼D(Ω) ℓ(g k m (h k m (x)), x) → ϵ→0 0as ϵ → 0, for any m, k. It remains to prove similar convergence results for the remaining terms in E R .η 2 E x∼D(Ω) ∥Jg k m,ϵ [h k m,ϵ (x)]∥ 2 F + ∥Jh k m,ϵ [x]∥ 2 F . We now show that Ω ∥Jg k m,ϵ [h k m,ϵ (x)]∥ 2 F dµ → Ω ∥Jg k m [h k m (x)]∥ 2F dµ via the dominated convergence theorem (DCT). First note that Ω(m, k) and Ω only differ by a set of µ-measure zero (the Voronoi boundaries S k m ), so we can equivalently proveΩ(m,k) ∥Jg k m,ϵ [h k m,ϵ (x)]∥ 2 F dµ → ϵ→0 Ω(m,k) ∥Jg k m [h k m (x)]∥ 2 F dµ.This allows us to avoid points x such that x or h k m (x) lie on Voronoi boundaries; these are problematic because if h k m (x) lies on a Voronoi boundary, then Jg k m [h k m (x)] is undefined. First note that as g k m , h k m are piecewise affine and the mollifiers are compactly supported, for any given x ∈ Ω(m, k), one can choose ϵ > 0 sufficiently small so thatg k m,ϵ [h k m,ϵ (x)] = g k m [h k m (x)] and hence ∥Jg k m,ϵ [h k m,ϵ (x)]∥ 2 F = ∥Jg k m [h k m (x)]∥ 2 F . In particular, for any x ∈ Ω(m, k), let ϵ 1 < d(x, S k m ) and ϵ 2 < d(h k m (x), S k m ); these can both be > 0 because d(x, S k m ) > 0, d(h k m (x), S k m ) > 0 for x ∈ Ω(m, k). Then define ϵ := min{ϵ 1 , ϵ 2 }. (21) shows why choosing ϵ 1 < d(x, S k m ) implies h k m,ϵ (x) = h k m (x); similar arguments hold show that ϵ 2 < d(h k m (x), S k m ) implies g k m,ϵ (h k m (x)) = g k m (h k m (x)). We then have g k m,ϵ [h k m,ϵ (x)] = g k m,ϵ [h k m (x)] = g k m [h k m,ϵ (x)] as desired. Hence for this choice of ϵ, Jg k m,ϵ [h k m,ϵ (x)] = Jg k m [h k m (x)] and consequently ∥Jg k m,ϵ [h k m,ϵ (x)]∥ 2 F = ∥Jg k m [h k m (x)]∥ 2 F . (Recall that Jg k m [h k m (x)] is well-defined for x ∈ Ω(m, k).) It follows that ∥Jg k m,ϵ [h k m,ϵ (x)]∥ 2 F → ∥Jg k m [h k m (x)]∥ 2 F pointwise on Ω(m, k)and pointwise µ-ae on Ω. Furthermore, another argument via Jensen's inequality shows that∥Jg k m,ϵ [h k m,ϵ (x)]∥ 2 F is dominated by the function A k m that is identically equal to sup x∈Ω(m,k) ∥Jg k m [h k m (x)]∥ 2 F ; this function is inte- grable because Ω is compact. The DCT then lets us conclude that ∥Jg k m,ϵ [h k m,ϵ (x)]∥ 2 F → ∥Jg k m [h k m (x)]∥ 2 F in L 1 and hence that Ω ∥Jg k m,ϵ [h k m,ϵ (x)]∥ 2 F dµ → ϵ→0 Ω ∥Jg k m [h k m (x)]∥ 2 F dµ.The same argument also allows us to conclude thatΩ ∥Jh k m,ϵ (x)]∥ 2 F dµ → ϵ→0 Ω ∥Jh k m (x)]∥ 2 F dµ.We have by now shown thatE R (g k m,ϵ , h k m,ϵ ) ≤ E R (g k m , h k m ) + ψ(ϵ; m, k)for some ψ(ϵ; m, k) → 0 as ϵ → 0. Combining this with our earlier results, we getE R (g k m,ϵ , h k m,ϵ ) ≤ E R (g k m , h k m ) + ψ(ϵ; m, k) ≤ inf f ∈C ∞ (Ω) E L (f ) + 1 m + 1 k + ψ(ϵ; m, k).As we can make 1 m + 1 k +ψ(ϵ; m, k) arbitrarily small by first choosing m, k sufficiently large and then choosing ϵ(m, k) > 0 to make ψ(ϵ; m, k) sufficiently small, we can finally conclude that (R) ≤ (L) as desired.This completes the proof of the theorem. ■ A.3 Proof of Theorem 3.2We will show thatσ 2 ∥Jf [x]∥ 2 F = E ϵ∼N (0,σ 2 I) ∥f (x + ϵ) -f (x)∥ 2 2 + O(σ 2 ).Since f : R n → R m is continuously differentiable, Taylor's theorem states that:f (x + ϵ) = f (x) + Jf [x]ϵ + R(x + ϵ),where∥R(x + ϵ)∥ 2 ∈ O(∥ϵ∥ 2 2 ).Rearranging, taking squared Euclidean norms, and expanding the square, we obtain:∥Jf [x]ϵ∥ 2 2 = ∥f (x + ϵ) -f (x) -R(x + ϵ)∥ 2 2 = ∥f (x + ϵ) -f (x)∥ 2 2 -2 • ⟨f (x + ϵ) -f (x), R(x + ϵ)⟩ + ∥R(x + ϵ)∥ 2 2 ∈O(∥ϵ∥ 4 2 ) ≤ ∥f (x + ϵ) -f (x)∥ 2 2 + 2 ∥R(x + ϵ)∥ 2 ∈O(∥ϵ∥ 2 2 ) • ∥f (x + ϵ) -f (x)∥ 2 ∈O(∥ϵ∥2) +O(∥ϵ∥ 4 2 ) = ∥f (x + ϵ) -f (x)∥ 2 2 + O(∥ϵ∥ 3 2 ) = ∥f (x + ϵ) -f (x)∥ 2 2 + O(∥ϵ∥ 2 2 ).Hutchinson's trace estimator implies that for any matrix A,∥A∥ 2 F = E ϵ∼N (0,I) [∥Aϵ∥ 2 2 ].In particular,σ 2 ∥Jf [x]∥ 2 F = E ϵ∼N (0,σ 2 I) ∥Jf [x]ϵ∥ 2 2 = E ϵ∼N (0,σ 2 I) ∥f (x + ϵ) -f (x)∥ 2 2 + O(∥ϵ∥ 3 2 ) = E ϵ∼N (0,σ 2 I) ∥f (x + ϵ) -f (x)∥ 2 2 + O(E∥ϵ∥ 2 2 =σ 2 n ) = E ϵ∼N (0,σ 2 I) ∥f (x + ϵ) -f (x)∥ 2 2 + O(σ 2 ),which completes the proof of the result. ■Let X ∈ R D×N be a low-rank matrix of clean data and Y = X + σ ϵ Z be a matrix of data corrupted by iid white noise Z. In this appendix, we show that the solution toEQUATIONcoincides with Gavish and Donoho [2017] 's optimal shrinker for the squared Frobenius norm loss when η is set to the noise variance σ 2 ϵ and as the "aspect ratio" β := d n → 0. A * is optimal for problem (28) iff 0 ∈ ∂ϕ(A * ). Using well-known results from convex optimization, this condition is equivalent to:EQUATIONFurthermore, the subgradient of the nuclear norm is:EQUATIONwhereA = U A Σ A V ⊤A is an SVD of A. We will show that the solution to (28) isEQUATIONwhere Y = U ΣV ⊤ is an SVD of the noisy data matrix, andEQUATION)The idea is to use the SVD Y = U ΣV ⊤ and the ansatz A * = U ΓU ⊤ to rewrite the LHS of the inclusion (29) as follows:1 N η (Y -A * Y )Y ⊤ = U 1 N η (I -Γ)Σ 2 U ⊤ .We then express the middle diagonal term as follows:1 N η (I -Γ)Σ 2 = I T + 1 N η (I -Γ)Σ 2 -I Twhere I T is the identity matrix with all columns ≥ T set to zero (i.e. an orthogonal projection matrix onto the first T coordinates). This then yieldsU 1 N η (I -Γ)Σ 2 U ⊤ = U T U ⊤ T + U 1 N η (I -Γ)Σ 2 -I T U ⊤ ,where U T = U I T is U with all columns ≥ T set to 0, and T is the first index such that σ T ≤ √ N η. As the optimal Γ in (32) sets all entries corresponding to singular values σ d ≤ √ N η to zero, we can in fact rewrite A * = (U I T )Γ(U I T ) ⊤ . It follows that U T U ⊤ T = U I T (U I T ) ⊤ is also of the formU A V ⊤A for a valid SVD of A * (the U I T can serve as both left-and right-singular vectors). We have therefore expressed the LHS of the inclusion (29) as:1 N η (Y -A * Y )Y ⊤ = U A V ⊤ A + W for U A V ⊤ A = U T U ⊤ T and W = U 1 N η (I -Γ)Σ 2 -I T U ⊤ .This W satisfies all of the conditions in (30).Furthermore, if we apply this optimal A * to the data matrix Y , we obtain A * Y = U ΓΣV ⊤ , whereEQUATIONIf we set η to be equal to the noise variance σ 2 ϵ , then this agrees exactly with the optimal shrinker for the case β := D N → 0 from Gavish and Donoho [2017] under the same noise model Y = X + σ ϵ Z. Architecture. In all ROF experiments, we parametrize f θ = g θ • h θ , where g θ and h θ are both two-layer MLPs with 100 hidden units. We apply a Fourier feature mapping [Tancik et al., 2020] to the input coordinates x before passing them through h θ . We use ELU activations in both neural nets and find the use of differentiable non-linearities to be crucial for obtaining accurate solutions.Training details. We train all neural models using the AdamW optimizer [Loshchilov and Hutter, 2019] at a learning rate of 10 -4 for 100,000 iterations with a batch size of 10,000. In the n = 2 case, we integrate over the box [-10, 10] 2 , and in the n = 5 case, we integrate over the box [-2, 2] 5 .We employ a warmup strategy for solving our problem (9). We first train our neural nets at η = 0.05 in the n = 2 case and η = 0.01 in the n = 5 case for 10,000 iterations, and then increase η by 0.05 and 0.01, respectively, each 10,000 iterations until we reach the desired value of η. We then continue training until we reach 100,000 total iterations.Each training run for (8) takes approximately 2 hours, and each training run for (9) takes approximately 45 minutes on a single V100 GPU.Architecture. Our architecture for all denoising models is based on the UNet implemented in the Github repository for Zhang et al. [2021] . Each model is of the form f = g • h, where h consists of the head, downsampling blocks, and body block of the Unet, and g consists of a repeated body block, the upsampling blocks, and the tail. We replace all ReLU activations with ELU but leave the remainder of the architecture unchanged.Training details. All neural models are trained on 288,049 images from the ImageNet Large Scale Visual Recognition Challenge 2012 training set [Russakovsky et al., 2015] which we randomly crop and rescale to 128 × 128. The code for loading and pre-processing this training data is borrowed from Rombach et al. [2021] .We train all neural models using the AdamW optimizer [Loshchilov and Hutter, 2019] for 2 epochs at a learning rate of 10 -4 then for a final epoch with learning rate 10 -5 . Each denoising model takes approximately 5 hours to train on a single V100 GPU.The training objective for our denoiser is (11) with η = σ. The training objective for the supervised denoiser is the usual MSE loss:inf f θ :R D →R D E x∼D(Ω) ϵ∼N (0,I) 1 2 ∥f θ (x + σϵ) -x∥ 2 2 ,where D(Ω) now denotes the empirical distribution over clean training images. The training objective for the Noise2Noise denoiser is:inf f θ :R D →R D E x∼D(Ω) ϵ1,ϵ2∼N (0,I) 1 2 ∥f θ (x + σϵ 1 ) -(x + σϵ 2 )∥ 2 2 .Note that this requires access to independent noisy copies of the same clean image during training.Evaluation details. We evaluate each denoiser by measuring their average peak signal-to-noise ratio (PSNR) in decibels (dB) on 100 randomly-drawn images from the ImageNet validation set, randomly cropped to 256 × 256. We corrupt each held-out image with Gaussian noise with the same standard deviation that the respective models were trained on (σ ∈ {1, 2}) and denoise them using each neural model along with BM3D [Dabov et al., 2007] , a popular classical baseline for unsupervised denoising.We use the β-VAE implementation from the AntixK PyTorch-VAE repo and use the default hyperparameters (in particular, we set β = 10) but set the latent dimension to 32, as we find that this yields more meaningful latent traversals. Training this β-VAE takes approximately 30 minutes on a single V100 GPU.We describe the architecture and training details for our regularized and unregularized autoencoder which we use to generate the latent traversals in Figures 6 and 7 . Training this autoencoder with the de Deterministic autoencoder architecture. Our autoencoder operates on 256 × 256 images from the CelebA dataset. To reduce the memory and compute costs of our autoencoder, we perform a discrete cosine transform (DCT) using the torch-dct package and keep only the first 80 DCT coefficients.We then pass these coefficients into our autoencoder.Our deterministic autoencoder consists of an encoder f θ followed by a decoder g ϕ . The encoder f θ is parametrized as a two-layer MLP with 10,000 hidden units; the latent space is 700-dimensional. The decoder g ϕ consists of a two-layer MLP with 10,000 hidden units and 3 * 80 * 80 = 19200 output dimensions, followed by an inverse DCT, and finally a UNet. We use the same UNet as in the denoising experiments described in Appendix B.2.Training details. We train our autoencoders with the following objective:EQUATIONwhere D(Ω) is the CelebA training set [Liu et al., 2015] .This is a standard deterministic autoencoder objective, with our regularizer approximating the Jacobian nuclear norm ∥Jf θ [x]∥ * of the encoder f θ .We train the unregularized autoencoder with η = 0, and the regularized autoencoder with η = 0.5. In both cases, we train on the CelebA training set for 4 epochs using the AdamW optimizer [Loshchilov and Hutter, 2019] with a learning rate of 10 -4 . Training these autoencoders takes approximately 4 hours each on a single V100 GPU.Generating latent traversals. To generate the latent traversals for our autoencoder, we draw a point x from the training set, compute the encoder Jacobian Jf θ [x], and take its SVD to obtain Jf θ [x] = U ΣV ⊤ . We then take the first 5 left-singular vectors (i.e. the first 5 columns of U ) and compute z = f θ (x) + αu d θ (x), where u d θ (x) denotes the d-th column of U . Here α denotes a scalar coefficient; it ranges over [-20000, 20000] for the unregularized autoencoder and [-2000, 2000] for the regularized autoencoder.To generate the latent traversals for the β-VAE, we encode the same training point x and replace the d-th latent coordinate with an equispaced traversal of [-3, 3] for d ∈ {1, 2, 11}; these are the first three meaningful latent traversals for this training point.38th Conference on Neural Information Processing Systems (NeurIPS 2024). 


Decision: This paper addresses the computational challenges in regularizing deep learning with the nuclear norm of the Jacobian matrix.  The key discovery is that that for functions parametrized as compositions $f = g \circ h$, the regularizer is equivalent to the average squared Frobenius norms of $Jg$ and $Jh$.  The latter two are then efficiently computed via a denoising-style approximation.  Empirical superiority is also shown in denoising and representation learning. 

The paper addresses a common challenge with a neat and effective solution.  The new proof in the rebuttal to Reviewer gXsR appears correct in my opinion, and I recommend the authors to rewrite the paper to reflect the changes.  Although the experiments can still be much improved as the reviewers suggested, it appears agreed that the merits of the paper are sufficient to warrant publication at the conference.

Reviews: review: paper_summary: paper_summary: The paper describes an elegant and very efficient numerical scheme for minimizing a regularization term taking the form of the nuclear norm of the Jacobian $\\|Jf [x]\\|_*$ of a function $f$ at an input $x$. The numerical scheme can be applied when the function is written as a composition of two functions $f=g\circ h$. The numerical scheme is supported by two theorems and experiments. The regularization term is supported by experiments.
strength_weakness: strength_weakness: For me, the main weaknesses of the article lie in the proof of the theorems (see below).

Another point that is unclear to me is due to the discrepancy between Theorem 3.1 and its application. Specifically, as it is currently written, on the right-hand side of (4), there is an infimum on all smooth functions  $h$ ad $g$. The "latent space" (the output space of $h$ and the input space of $g$) is not specified, and my understanding is that the infimum is over the union of all possible "latent spaces". The infimum over this union can be much smaller than the infimum for a given "latent space"... and, in practice, only one latent space is considered. This point should be clarified. Perhaps the theorem can be adapted to allow the freedom to choose the latent space.

Similarly, in practice, neural networks (of fixed architecture) are not able to reach all the smooth functions. This could have an impact on the two infimums of (4). This limitation should at least be stated clearly. 

Finally, throughout the article, the reader would welcome hints on the choice of $g$ and $h$. In neural networks, we have $h_L \\circ...\\circ h_2 \\circ h_1$ and the reader doesn't know whether the choice of $h$ and $g$ resumes to choosing a layer separating them? or whether the method is extended to more than one composition in some other way? 

Robustness to adversarial attacks seems to be a natural advantage of the regularization term. The authors may wish to mention this, along with other perspectives they deem relevant, in the conclusion.

In (3) and all the developments, the authors write $l(f(x),x)$. It is much more common in machine learning to consider  $l(f(x),y)$, for an input/output pair $(x,y)$. The article would really gain to be extended to this setting.

In Theorem 3.2: I think it is a little o: $o(\\sigma^2)$, a quantity that goes to $0$ faster than $\\sigma^2$.

{\bf Comments on the appendices:}

Line 455: Please define AM-GM.

Line 457: Please remove 'compact'. It is the usual SVD.

Line 492: A sentence like 'Consider $z\\in\\Omega$' would be welcome.

In (14), last term: I think it is $f_m(z) + R_m^z(x)$, not  $f_m(x) + R_m^z(x)$. This leads to several similar changes which (I think) do not have serious consequences. The changes are in (16), twice in line 498, once in line 509, and once below line 514.

Line 503: Please replace 'These functions....' with 'The composition of these functions'.

Line 509: I think the calculation would be simpler if you replace $ \\|Jf_m[z_i] - Jf_m[x]\\|_{*} $ 

by the absolute value of the difference between $  \\| Jf_m[z_i]\\|_{*} $ 

and $ \\| Jf_m[x]\\|_{*}  $.

Calculus below Line 514: The layout is odd. Also, the calculation would gain clarity if you write the sum in $i$. and upper-bound each term.

Above Line 531 (and throughout the proof): Please write $\\epsilon \\rightarrow 0$ under the arrow in $\\rightarrow 0$.

Line 533: Please define RHS

Line 535: The fact that $\\|h^k_{m,\\epsilon} - h^k_m\\|_{L^1(\\Omega,\\mu)} \rightarrow 0$ is not clear to me. For instance, $\mu$ might involve a weighted Dirac mass at a point where $h^k_m$ is discontinuous. Although this paragraph is a sort of illustration, it is preferable to avoid saying something that might be false. Finally, at some other location in the proof, I think you need $\mu$ to be absolutely continuous.

Line 542: Please precise 'complement in $\\Omega$'.

Line 545: It is $g^k_{m,\\epsilon}$, not $g^k_{m,e}$.

Equation below Line 546: It is $B(h^k_m(x), \\epsilon)$ , not $B(h^k_m(x))$.

**A major issue:** From lines 543 to 549, the authors try to prove an inequality which, I think, cannot hold. The problematic intermediate step is in the inequality below line 546. It is

$\\|Jg^k_{m,\\epsilon} [H_m^k(x)]\\|_F \\leq \\sup_{y\\in h_m^k(\\Omega)  \\|Jg_m^k[y]\\|_F }.$

It cannot hold because $g^k_{m,\\epsilon}$ is a smooth approximation of $g^k_m$ which is typically discontinuous. It is possible to find $\epsilon$ close to $0$ and points $y$ such that $\\|Jg^k_{m,\\epsilon} [y]\\|_F$ is arbitrary large. To me, there is no guarantee that $h^k_m(x)$ avoids such points. Said differently, to me, the first inequality below Line 546 does not hold when $B(h^k_m(x),\\epsilon) \\cup V_j\\neq \\emptyset$, where $V_j$ is a Voronoi cell such that $h^k_m(x)\\not\\in V_j$, which your hypotheses do not exclude. By the way, you might want to provide the formula for the mollification and detail the calculations similar to the first inequality below Line 546.

Above line 550: It is $L^1(\\Omega,\\mu)$, not $L^1(\\Omega)$

Line 555: You state: 'As $h^k_{m,\\epsilon} \\rightarrow h^k_{m}$, $\\mu(\\Omega_1(m,k,\\epsilon)) \rightarrow 0$.'. It would be useful to state in which sense $h^k_{m,\\epsilon} \\rightarrow h^k_{m}$ and to prove the details of the arguments guaranteeing that $\\mu(\\Omega_1(m,k,\epsilon)) \\rightarrow 0$. As already said, I suspect you need $\\mu$ to be absolutely continuous for the conclusion to hold. If I am correct, this hypothesis should appear in Theorem 3.1.

**A major issue:** Line 568, you state $\\int_\\Omega \\|Jg^k_{m,\\epsilon} [h^k_{m,\\epsilon} (x)]  \\|^2_F d\\mu \\rightarrow \\int_\\Omega \\|Jg^k_{m} [h^k_{m} (x)]  \\|^2_F d\\mu$. The sense of the term on the right of the equality is not clear to me since $g^k_{m}$ is piecewise linear, and might even be discontinuous. Even worth, I think nothing excludes $h^k_m$ to be constant on one of the Voronoi cell $V$. If this happens and if the constant turns out to be a point such that $Jg^k_{m}$ does not exist, the problem occurs on $V$ and we generally have $\\mu(V) \\neq 0$. Concerning the term on the left of equality sign, I fear you encounter problems similar to those you may be familiar with since you mention the Total Variation in your article. The term on the left takes the jumps into account and might even go to infinity since you have a square. Looking at the proof of this statement, we find a possible cause for the mistake in lines 569-570. There, you state `...for any given $x\\in\\Omega$ one can choose $\\epsilon>0$ sufficiently small so that $g^k_{m,\\epsilon}(h^k_{m,\\epsilon} (x)) = g^k_m(h^k_{m} (x))$...'. This is not true if $x$ is on the boundary of the Voronoi cells defining $h^k_m$ or is such that $h^k_{m} (x)$ is the boundary of two Voronoi cells defining $g^k_m$.

Appendix A.3: You always use the (big) $O$ but sometimes you have to write the (little) $o$. You need to distinguish between the two notations.
questions: questions: SVS and SVT are often interpreted as proximal operators. Can you please mention it in your introduction?
limitations: limitations: The main limitation is that there are gaps in the proof of the main theorem.
paper_summary: paper_summary: The authors present an efficient method for regularizing the Jacobian of deep networks such that it is low-rank. This work is motivated by the fact that penalizing the Jacobian by the nuclear norm regularization is in general a computationally difficult task, as it needs to (i) actually compute the Jacobian and (ii) take the SVD of a large matrix. The authors' proposed method comes from the observation that the nuclear norm of a matrix can be computed by a non-convex optimization problem, which I believe is somewhat commonly considered in the matrix factorization literature. Then, they propose their method for estimating the Jacobian Frobenius norm, which is equivalent to computing the nuclear norm (roughly speaking).
strength_weakness: strength_weakness: - I think the main weakness of the paper is its experimental section. While I think this method could be a good starting point for other future papers, the current experiments don't really sell the effectiveness of the method. It seems that one of the main benefits of this method as shown in the experimental results in Table 1 is that the proposed denoiser is almost as good as supervised denoiser, despite having no corresponding clean images for training. But isn't N2N also unsupervised in the sense that it doesn't need clean pairs of images? It doesn't seem to have impressive performance gains over N2N or BM3D.
- Building upon the previous weakness, I wonder if there is a way of showing the effectiveness of this method through other means -- for example computational efficiency. This paper was motivated by the fact that it can circumvent costly Jacobian + SVD computations. Could the authors show with a small scale experiment showing the computational gains? I think that would strengthen the paper.
questions: questions: I do not have any specific questions besides the ones in the weaknesses.
limitations: limitations: The limitations are listed in the weaknesses section.
paper_summary: paper_summary: The paper proposes a computationally tractable method to induce low-rankness of a neural network's Jacobian. The method essentially generalizes the max-norm from Renni and Srebro, 2005, to more general compositions of functions, as it is common in neural networks. The method is made computationally efficient by estimating the Frobenius norm of the Jacobian. Experiments for denoising and representation learning are used to validate the efficacy of the method.
strength_weakness: strength_weakness: While the preliminaries in Section 3.1 are easy to follow, large parts of Section 3.2, which contains the main contribution, would be easier to follow if the authors would give a concrete example early on for the functions f and g. An illustrative example would make the idea more accessible.

One of the key arguments of the paper is that the method scales to high-dimensional problems and very large neural networks. The application examples, however, consider relatively small and simple problems with relatively simple architectures (for both the denoising of not-so-large images and autoencoder-based representation learning of images tasks). It is therefore not demonstrated that the method indeed scales to high-dimensional problems and complex neural networks.

The impact of estimating the Frobenius norm of the Jacobian (as detailed in Section 3.3) on the performance of applications was not studied. The authors claim that a single draw of eps is sufficient to evaluate (6), but it is unclear to me what the impact of this choice is. 

It might be good if the authors could quantify the complexity of the proposed method, when used during training (either in run-time or in the number of operations/multiplications). This could also provide insight into the scalability to large-dimensional problems with very large neural networks.
questions: questions: How are hyperparameters (e.g., the regularization factor eta) in the provided applications set?

Denoising and representation learning can be accomplished with other means as well, and existing state-of-the-art methods are likely to outperform the proposed approach. Is there any application that would uniquely benefit from the proposed regularizer, i.e., in which no other existing method can be used?
limitations: limitations: I do not see a specific limitation that was not discussed.
paper_summary: paper_summary: The paper proposes a method to regularize the nuclear norm of the Jacobian of a function, e.g., one that represents a neural network. This method builds on prior art and includes the authors’ novel contribution as follows: The authors reference prior art for an equivalent problem formulation of nuclear norm regularization that avoids the computation of SVD. This equivalence also enables the authors’ novel result, which is that the nuclear norm of the Jacobian $Jf$ of a composite function $f=g \circ h $ is equal to the average of the squared Frobenius norms of $Jg$ and $Jh$; this result makes the proposed method perfectly apt for neural network training. The authors then use an approximation to the Frobenius norm of the Jacobians to avoid explicitly calculating a large Jacobian matrix, which significantly reduces computational cost and storage. The method is validated with a nuclear norm-regularized problem whose closed form solution is known. Finally, the efficacy of the method is shown on two applications: unsupervised denoising and representation learning.
strength_weakness: strength_weakness: Major comment:
- In the representation learning application, we only see the method’s efficacy on a single image. Could there be a way to quantify the performance of the method over a whole test dataset? 

Minor comments: 
- I found the following two references that also compute nuclear norm without the SVD. Could the authors either mention these references in their manuscript or clarify why these aren’t relevant?
[1] https://icml.cc/Conferences/2010/papers/196.pdf
[2] https://www.ijcai.org/proceedings/2017/0436.pdf

- Figure 1 has only the color legend. Could the authors confirm (and explicitly state on their manuscript) that the x- and y-axes correspond to each coordinate of the input x?
- In the caption of Figure 3, the authors say that “As predicted by Theorem 3.1, both problems converge to the same objective value.“ Could they use a less strong claim here? Possibly “nearly identical” as they said in the main text; otherwise, the gap in 3(c) is confusing.
- Could the authors reformat the subsection titles in Section 5 so that it is clear in a glance that “Unsupervised denoising” and “Representation learning” are the two application examples, and “Singular value shrinkage” and “Experiments” are under “Unsupervised denoising”?
- The PSNR values in Table 1 are based on averaging over only 100 images. Is there a reason why more images were not used? Could the authors repeat the experiment by averaging over more images?
- It is slightly confusing that the order of the methods in the Table 1 and Figure 4 do not match each other. Could the authors follow the same order?
questions: questions: Please see Weaknesses.
limitations: limitations: The authors briefly comment on the limitation of their method regarding the approximation of the squared Frobenius norm of the Jacobian in Conclusions. I think this is sufficient.

Comments: comment: We thank Reviewer gXsR for reading our rebuttal. We are confident that our updated proof is now correct, and we believe the details we have provided in our rebuttal are sufficient to resolve the specific concerns raised by Reviewer gXsR under the "major issue" headings.
I have read the author's rebuttal but, in my opinion, the proof needs to be re-read in detail and the article needs another round of review. It is not possible, from an article and a rebuttal, to check that the proof is correct. For this reason, I will not change my rating. To avoid errors in the futur, I recommend writing down the details of the mollification and being careful when swapping limits and integrals.

Best regards,
Thank you for answering my questions along with the additional experiments; I have raised my score accordingly.
Thank you for clearly addressing my comments!