Nuclear Norm Regularization for Deep Learning

Thank you for your thoughtful review of our manuscript. We are glad you appreciate the originality and significance of our method for Jacobian nuclear norm regularization.

"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?"

We depict latent traversals in our autoencoder's latent space for a single non-cherry-picked image due to space constraints in the manuscript. We have included a few more examples in our rebuttal PDF attached to the global response, and we would be happy to include these in an appendix in the camera-ready. Whether our autoencoder recovers semantically meaningful directions of variation in latent space is ultimately a subjective judgment, and we do not believe that readers would gain additional insight regarding our method's performance on this task by including quantitative metrics.

"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?"

We would be pleased to add these references to the related work section. However, note that both of these papers propose methods for nuclear norm regularization in matrix learning problems, where one seeks to learn a single matrix $A \in \mathbb{R}^{m\times n}$. In contrast, our work generalizes the efficient method of Rennie and Srebro [2005] to non-linear learning problems, where the appropriate analog to penalizing the nuclear norm of a matrix is penalizing the nuclear norm of the Jacobian of the function being learned.

"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?"

That is correct. We will clarify this in the camera-ready.

"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?"

The standard test sets in the denoising literature are quite small; for example, our other test set "CBSD68" contains 68 images. We built our Imagenet test set using 100 random images to approximately match test set sizes that are common in the denoising literature. We do not expect that the results of our comparison would be significantly different if we used a larger test set.

We would be happy to incorporate the rest of your suggested changes to our paper's formatting in the camera-ready.

We hope this answers your questions and would be pleased to continue this conversation in the author-reviewer discussion period.

Thank you for your thoughtful review of our manuscript. We are glad that you appreciate our technical contributions and the practical utility of our method.

"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."

In our denoising experiment, we apply our regularizer to a function of the form $f_\theta : \mathbb{R}^d \rightarrow \mathbb{R}^d$, where $d = 3 \times 256 \times 256$ since it operates on $256 \times 256$ RGB images. In this case, $Jf_\theta[x]$ is a $(3 \times 256 \times 256) \times (3 \times 256 \times 256) = 196608 \times 196608$ matrix, which occupies over 154 GB of memory if stored in float32 format and therefore cannot be stored on most GPUs. As our regularizer does not require any Jacobian computations, it can be successfully applied to this problem. In contrast, naive Jacobian nuclear norm regularization would be intractable due to the need to compute and take the SVD of the $196608 \times 196608$ denoiser Jacobians.

"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."

We provide the formula for our regularizer in Equation (6) in Section 3.3 of our manuscript. As $f(x) = g(h(x))$, one must evaluate the $h(x)$ and $g(h(x))$ terms to train $f$ regardless of the training objective. Our regularizer requires additionally computing $h(x+\epsilon)$ and $g(h(x) + \epsilon)$, so if one computes the regularizer with $N$ noise samples, our regularizer requires an additional $2N$ function evaluations. As we used $N=1$ noise sample to compute our regularizer in our denoising and autoencoding experiments, the marginal cost of our regularizer was 2 function evaluations per iteration.

This compares favorably to computing the Jacobian, which requires $O(d)$ function evaluations for a function $f : \mathbb{R}^d \rightarrow \mathbb{R}^d$, and to the $O(d^3)$ SVD computation which is then required to compute the gradient of the Jacobian nuclear norm. Furthermore, as noted above, merely storing the Jacobian matrix is intractable for high-dimensional problems such as denoising.

"How are hyperparameters (e.g., the regularization factor eta) in the provided applications set?"

In the denoising experiments, we set $\eta = \sigma^2$, where $\sigma^2$ is the noise variance. (In our current manuscript, we have indicated that we set $\eta = \sigma$ in these experiments; this is a typo that we will correct in the camera-ready.) This choice of $\eta$ is motivated by results on optimal singular value shrinkage from Gavish and Donoho [2017], whose optimal shrinker solves a special case of our proposed denoising problem (11). (See our section "Singular value shrinkage" from line 230 onwards for details.) We set $\eta$ empirically in the autoencoder experiment via grid search.

"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?"

Our primary contribution is a tractable and well-grounded method for Jacobian nuclear norm regularization that can be applied to high-dimensional deep learning problems. There are few problems for which some given regularizer is strictly necessary and no other approach can be used, and our regularizer is no exception to this principle. Our goal in our experiments was to demonstrate the practical value of our proposed regularizer in high-dimensional learning problems by showing that it performs well in two tasks of interest to the machine learning community. The efficiency and simplicity of our method will enable the community to build on our work and discover new applications for Jacobian nuclear norm regularization.

We hope this answers your questions and would be pleased to continue this conversation during the author-reviewer discussion period.

Thank you for your thorough review of our paper. We would be pleased to make the edits you have suggested for clarity and fix the typos you have pointed out in the camera-ready version of our paper. We have updated our proof of Theorem 3.1 to address your major issues; as we cannot upload revised manuscripts during the rebuttal period, we summarize our revised arguments below.

1. Our proof requires $\mu$ to be absolutely continuous wrt the Lebesgue measure $\lambda$ to transfer results of the form $\lambda(E) = 0$ or $\lambda(E_n) \rightarrow 0$ for sets $E_n, E \subseteq \Omega$ to $\mu(E) = 0, \mu(E_n) \rightarrow 0$, resp. Thank you for pointing this out; we will add this hypothesis to Theorem 3.1 in the camera-ready.

2. For any Voronoi partition $V_i, i=1,...,N(k)$ and corresponding functions $h_m^k, g_m^k$ that we construct in our proof, the preimage of the union of Voronoi boundaries (which we call $S_m^k$) under $h_m^k$ is a set of Lebesgue measure zero. If $\mu \ll \lambda$, then the preimage has $\mu$-measure zero as well. Our reasoning is as follows:

In lines 507-514, we define $g_m^k, h_m^k$ as piecewise affine functions such that $g_m^k(x) := g_m^{z_i}(x)$ and $h_m^k(x) := h_m^{z_i}(x)$ for all $x \in \textrm{int}(V_i)$. The affine functions $g_m^{z_i}, h_m^{z_i}$ have the key property that $\frac{\eta}{2}\left( \|Jg_m^{z_i}[h_m^{z_i}(x)]\|_F^2 + \|Jh_m^{z_i}[x]\|_F^2 \right) = \eta \|J f_m[z_i]\|\_*.$

At any $x$ on the interior of a Voronoi cell $V_i$, $g_m^k$ and $h_m^k$ are equal to $g_m^{z_i}, h_m^{z_i}$, resp, so by the key property above, $\frac{\eta}{2}\left( \|Jg_m^k[h_m^k(x)]\|_F^2 + \|Jh_m^k[x]\|_F^2 \right) = \eta \|J f_m[z_i]\|\_* < + \infty$. In particular, $\|Jg_m^k[h_m^k(x)]\|_F^2$ and $\|Jh_m^k[x]\|_F^2$ must both be well-defined and finite at this $x$. This can only happen if $h_m^k(x)$ lies on the interior of a Voronoi cell, as otherwise $Jg_m^k[h_m^k(x)]$ would be undefined. Hence $h_m^k$ maps the interiors of Voronoi cells to the interiors of Voronoi cells; the contrapositive is that if $h_m^k(x)$ lies on a Voronoi boundary, then $x$ also lies on a Voronoi boundary. Consequently, the preimage of the Voronoi boundaries under $h_m^k$ is a set of Lebesgue measure zero. If $\mu \ll \lambda$, this is also a set of $\mu$-measure zero.

Using this fact, we can replace all integrals over $\Omega$ from line 525 onwards with integrals over $\Omega(m,k) := \Omega \setminus S_m^k$ without changing their value. In particular, $Jg_m^k[h_m^k(x)]$ is well-defined $\mu$-ae, which should resolve your concerns re: line 568, for which you note "I think nothing excludes $h_m^k$ 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_m^k$ does not exist, the problem occurs on $V$ and we generally have $\mu(V) = 0$." It should also resolve your concerns re: lines 569-570, for which you state that one cannot always find $\epsilon>0$ so that $g_{m,\epsilon}^k[h_{m,\epsilon}^k(x)] = g_m^k[h_m^k(x)]$. One can in fact find such $\epsilon$ for $\mu$-almost all $x$, which is sufficient to apply the dominated convergence theorem.

3. We use an alternative approach to show that $\|g_{m,\epsilon}^k \circ h_{m,\epsilon}^k - g_{m,\epsilon}^k \circ h_{m}^k\|_{L^1(\Omega,\mu)} \rightarrow 0$ that avoids the issue you highlight re: lines 543-549. We employ a different decomposition of $\Omega \setminus S_m^k$ into good and bad sets:

• The good set $\Omega_0(m,k,\epsilon) \subseteq \Omega(m,k)$ is the set of $x \in \Omega(m,k)$ such that $d(x, S_m^k) > \epsilon$. ($d(p,S)$ denotes the distance from point $p$ from set $S$.)
• The bad set $\Omega_1(m,k,\epsilon)$ is $\Omega(m,k) \setminus \Omega_0(m,k,\epsilon) =$x \in \Omega(m,k) : d(x, S_m^k) \leq \epsilon $$.

For all $x \in \Omega_0(m,k,\epsilon)$, $h_{m,\epsilon}^k(x) = h_m^k(x)$ because $d(x, S_m^k) > \epsilon$ and we employ the standard mollifier supported on $B(0,\epsilon)$. Consequently, $g_{m,\epsilon}^k(h_{m,\epsilon}^k(x)) = g_{m,\epsilon}^k(h_m^k(x))$ and therefore $\|g_{m,\epsilon}^k(h_{m,\epsilon}^k(x)) - g_{m,\epsilon}^k(h_{m}^k(x)) \|\_2 = 0$, so $\int_{\Omega_0(m,k,\epsilon)} \|g_{m,\epsilon}^k(h_{m,\epsilon}^k(x)) - g_{m,\epsilon}^k(h_{m}^k(x)) \|_2 d\mu = 0$. We no longer need the reasoning in lines 543-549 to prove this integral converges to 0.

We now address $\int_{\Omega_1(m,k,\epsilon)} \|g_{m,\epsilon}^k(h_{m,\epsilon}^k(x)) - g_{m,\epsilon}^k(h_{m}^k(x)) \|\_2 d\mu$. We employ the same bound on the integrand as under line 552 for the new bad set $\Omega_1(m,k,\epsilon)$. To see that $\mu(\Omega_1(m,k,\epsilon)) \rightarrow 0$, note that $\lambda(\Omega_1(m,k,\epsilon)) \rightarrow 0$, as $\Omega_1(m,k,\epsilon)$ is a union of cylinders of radius $\epsilon$ centered at the Voronoi boundaries $S_m^k$, which have Lebesgue measure zero. Under our new assumption $\mu \ll \lambda$, it follows that $\mu(\Omega_1(m,k,\epsilon)) \rightarrow 0$.

These results jointly show that $\|g_{m,\epsilon}^k \circ h_{m,\epsilon}^k - g_{m,\epsilon}^k \circ h_{m}^k\|_{L^1(\Omega,\mu)}$ while avoiding the issue you highlight re: lines 543-549.

We hope this resolves your concerns regarding our proof. If you are satisfied with our response, we respectfully ask that you raise your score for our submission. Otherwise, we would be please to continue this discussion during the author-reviewer discussion period.

Thank you for your thoughtful review of our manuscript. We are glad that you appreciate our theoretical contributions and view our method as well-motivated. In this rebuttal, we will address the questions in your review. If you are satisfied with our answers, we respectfully ask that you raise your score for our submission.

"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."

As you note, our method performs nearly as well as a supervised denoiser and comparably to Noise2Noise (N2N), despite being trained exclusively on highly corrupted data without access to clean images. While N2N is also an unsupervised denoising method, training a denoiser using N2N requires several noisy samples of each clean image. Given a clean image $x$, these samples are of the form $x + \epsilon_i$, where the $\epsilon_i$ are distinct realizations of zero-mean noise. Such data is typically unavailable if one lacks access to the clean images, rendering N2N impractical for real-world applications. In contrast, our method requires only a single realization of each noisy image and can hence be applied to arbitrary datasets of noisy images, which can be easily obtained in the wild.

"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."

In our rebuttal PDF attached to the global response, we have included a figure comparing time per training step at batch size 1 for the denoising problem (11) using our regularizer and a naive Pytorch implementation of the Jacobian nuclear norm. We experiment with images drawn from Imagenet downsampled to sizes $S \times S$ for $S \in \{8,16,32,64\}$; our V100 GPU's memory overflows for $S \geq 128$. Whereas time per training step with the naive nuclear norm implementation rises to nearly 129 seconds per iteration for $64 \times 64$ images, each training step with our regularizer takes under 120 milliseconds. We additionally highlight that a key advantage of our regularizer is that it is tractable for problem sizes where simply computing the model Jacobian is infeasible. For example, in our denoising experiments, we trained our denoiser on $256 \times 256$ RGB images. The model Jacobian in this case is a $(3 \times 256 \times 256) \times (3 \times 256 \times 256) = 196608 \times 196608$ matrix, which occupies over 154 GB of memory if stored in float32 format and therefore cannot be stored on most GPUs. As our regularizer does not require any Jacobian computations, it can be applied to problems such as denoising where Jacobian computations are prohibitive.

We hope this resolves your concerns. If you are satisfied with our response, we respectfully ask that you raise your score for our submission. Otherwise, we would be please to continue this discussion during the author-reviewer discussion period.