Neural Eigenfunctions Are Structured Representation Learners

We appreciate the positive comments, and we are encouraged by the recognition of our presentation and contribution. We address the detailed comments below.

Regarding code

We clarify that our graph representation learning objective Eq. (12) is different from the SSL objective in Section 3. Specifically, instead of utilizing the argumentation kernel, we employed the graph adjacency kernel. As a result, data augmentation (represented as $X^+$) is not needed.

We also clarify that we did not square the first term in Eq. 12. The code for computing this term is in L75 of code/graph/main_products.py.

Thank you for taking the time to review our paper. We are glad that you found our reasoning straightforward and our results compelling. We have provided responses to the detailed questions below.

Regarding enabling/disabling stop_grad

We first apologize for typos in Eq 3 of the original manuscript and have revised it accordingly. We clarify that the second term $R_{i,j}^2$ in Eq 3 is asymmetric—inducing updates for only $\\psi_j$ instead of $\\psi_i, i< j$. Therefore, when optimizing the $k$ problems simultaneously using a shared neural network $\\psi$, we obtain the “surrogate” loss in Eq 8, where the stop_grad operation implements such asymmetry. In the presence of the stop_grad, the learning strictly follows Eq 3, and the learning outcomes are ordered eigenfunctions, so we obtain structured representations where choosing elements with small indices corresponds to choosing eigenfunctions with large eigenvalues (the “principal” components).

We remove the stop_grad operation for linear probe evaluation. This is because the linear probe does not take advantage of ordered representations (as shown in HaoChen et al., 2021, Lemma 3.1), while such a replacement may actually ease optimization because it relaxes the ordering constraints. Yet, after doing so, the ordered structure of the learned representation vanishes, so to perform image retrieval, we have to randomly select elements for an equal-length comparison and perform multiple runs.

Similarity in formulation to Barlow Twins while numerical comparison focused on a few experiments.

As explained in the paragraph Connection to Barlow Twins, the objectives of the two approaches are similar but the gradients and optimal solutions are different. The Neural Eigenmap objective is theoretically grounded and is guaranteed to recover the eigenfunctions span, while Barlow Twins can be seen as a less-principled approximation to this objective. Additionally, Barlow Twins lacks the ability to generate ordered representations and handle graph representation learning. For the empirical evaluation, we have shown statistically significant improvements over Barlow Twins across three tasks (ImageNet linear probe, COCO detection, and instance segmentation), a variety of batch sizes and projector dimensions. We believe these gains are meaningful and kindly ask the reviewer to reconsider their evaluation of these results.

Define BN layers and other typos

Thanks. We have revised the paper accordingly.

Can the optimization in Eq 3 be connected to the original definition of the eigenfunction from Eq 1?

Yes, as proved in the original NeuralEF paper (Deng et al. 2022, Theorem 1), the optima of the optimization problems in Eq 3 are the principal eigenfunctions, satisfying Eq 1.

How does the ordered structure arise in NeuralEF? Orthogonality requirement not present in (8)

We clarify that although the orthogonality requirement is not explicitly enforced in Eq. 8 (which is derived from the NeuralEF problem in Eq. 3), the orthogonality is automatically achieved at the optimum of the problem.

Intuitively speaking, the first term in Eq 3 corresponds to finding the principal eigenvectors by maximizing the Rayleigh quotient. When estimating multiple eigenvectors simultaneously, we also need to ensure that the $j$-th eigenvector is orthogonal to the $i$-th eigenvector for all $i < j$. The second term in Eq 3 is a penalty added to achieve this, as proven in the original NeuralEF paper (Deng et al., 2022).

$\\gamma$ should be $\\lambda$

Thanks. Yes, it is a typo. We have updated the manuscript.

Why Eq 9 works better than Eq 10

We believe Equation 9 works better than Equation 10 because of its theoretical underpinnings, ensuring the recovery of the eigenfunction span at the optimum. Conversely, we view Equation 10 as an “incorrect way” to implement the same idea.

The scalability of the proposed approach vs. that of Johnson et al. (2022)?

Our approach is as scalable as regular self-supervised learning methods—we demonstrate this on ImageNet-scale experiments and a graph representation learning benchmark with more than one million nodes, while Johnson et al. (2022) rely on the nonparametric kernel PCA, which has a cubic complexity w.r.t. the number of training points. Their experiments were conducted on 2D-synthetic datasets and MNIST.

Section 6.2 first sentence refers to Section 6.2

We apologize for the typo. We have revised the manuscript.

We appreciate the positive comments, and we are encouraged by the recognition of our contributions compared to existing works and the empirical effectiveness of the proposal. We address the detailed concerns below.

Which kernel to predefine

The purpose of our Section 3 and 4 is to address the question. In Section 3, along with the accompanying experiments, we showed strong evidence that the augmentation kernel is an excellent option for visual representation learning. Furthermore, in Section 4 and 6.3, we demonstrated good performance can be achieved by graph adjacency kernels in graph representation learning tasks.

Regarding the linear probe experiments for unsupervised representation learning

Why have the authors reproduced Barlow Twins themselves? This is because the linear probe results in the paper of Barlow Twins correspond to 1000-epoch training, while we perform 100-epoch training (as done in the paper of SCL) for results in Table 1 due to limited resources.

Can you make any inferential comments on the link between batch size, number of epochs, and projector dimension? Generally, there are no trade-offs between these three factors. For our method, batch size and projector dimension need both to be large enough for good linear probe performance, regardless of the number of epochs. A larger batch size contributes to a more reliable Monte Carlo estimation of the expectation in Eq 7, and a higher projector dimension translates to more eigenfunctions involved in the representation. For other self-supervised methods, different batch sizes and projector dimensions are preferred by different methods, e.g., contrastive learning methods without memory buffers may need a large batch size while others may not; SimCLR is robust to projector dimension, while Barlow Twins needs large projector dimensions.

Why are the results of Barlow Twins of a batch size of 2048 reported? We perform a new set of experiments for Barlow Twins using a batch size of 1024 and projector dimension of 8192, and obtain 65.5±0.1% top-1 linear probe accuracy on ImageNet, slightly weaker than using a batch size of 2048 for the projector dimension of 8192. The gap can stem from that we have not carefully tuned the learning rate for the LARS optimizer, but it is not significant, which echoes the results in Barlow Twins’ paper that it is relatively robust against batch sizes from 512 to 2048.

Connection between PCA and the proposed method

Your understanding is totally correct. We use such an analogy to intuitively explain why our learned representation is structured and ordered.

Why replacing $\\hat{\\psi}$ by $\\psi$ does not affect the optimal classifier for linear probe

Such a replacement makes the learned representation converge to the output of an arbitrary span of eigenfunctions. It amounts to applying a right transformation on the original eigenmaps. As proven by Lemma 3.1 of HaoChen et al., (2021), the effect of such a transformation can be easily eliminated by a trivial manipulation of the weight matrix of the linear classifier.

When does feature importance matter?

All benchmarks that use linear probe evaluations (Sections 6.2 and 6.3) ignore the significance of individual features. However, the image retrieval experiment in Section 6.1 utilizes the importance of features as they affect the tradeoff between cost and quality. By truncating less important features, our method can achieve similar retrieval quality using a much shorter code length than other methods that lack feature ordering.

Computational complexity for graph data

We first clarify that we did not state that graph NNs have a computational complexity of $O(n^3)$. Instead, it refers to classic nonparametric graph embedding methods like Laplacian Eigenmaps that involve matrix decomposition, which do have that complexity. Due to stochastic training, the computational complexity for our training is linear w.r.t. the number of data points $n$. During inference, when computing the embedding of a single node, our method simply requires forward propagation of a neural network with the features of that specific node as input, while graph NNs require aggregating information from the neighborhood on graphs. This neighborhood can be large, and hence the cost may be high.

We appreciate the constructive feedback. Below, we address the detailed concerns.

Clarification regarding Johnson et al. (2022)

In reference to the extended conference version [1] of Johnson et al.'s (2022) paper, we see that they also investigated using NeuralEF as an alternative to their main kernel PCA approach. However, we want to emphasize that this should be considered as concurrent to our work, as evident from our submission history. We will make this more clear in the paper. Additionally, it's important to note that the two papers have distinct focuses. [1] primarily investigates the optimality of representations obtained through kernel PCA. It only mentioned NeuralEF as an alternative and tested it in synthetic tasks. In contrast, our work extends the application of NeuralEF to larger-scale problems such as ImageNet-scale SSL and graph representation learning. Notably, [1] did not discuss the benefit of ordered representation while we clearly demonstrated it in the image retrieval tasks.

[1] Johnson, D. D., El Hanchi, A., & Maddison, C. J. Contrastive Learning Can Find An Optimal Basis For Approximately View-Invariant Functions. In International Conference on Learning Representations, 2023.

When the $\\hat{\\psi}$ in the second term of Eq. 8 should be used and when not

We recommend the use of $\\hat{\\psi}$ when downstream tasks require an ordered representation. However, in cases where the order of features has no impact, it is preferable to use $\\psi$. In our paper, we employed two distinct tasks to demonstrate this concept. In Section 6.2, when evaluating visual representation learning through linear probe, we used $\\psi$ in the second term. This choice was due to the fact that the optimal classifier for this task is unaffected by the ordering of eigenfunctions (as shown in HaoChen et al., 2021, Lemma 3.1). Conversely, for the image retrieval task discussed in Section 6.1, we employed $\\hat{\\psi}$ in Eq. 8. This allowed us to enforce feature ordering, enabling us to achieve the best quality-cost tradeoff (Figure 2) by discarding less important features.

Similarity to Barlow Twins

As explained in the paragraph Connection to Barlow Twins, the objectives of the two approaches are similar but the gradients and optimal solutions are different. The Neural Eigenmap objective is theoretically grounded and is guaranteed to recover the eigenfunctions span, while Barlow Twins can be seen as a less-principled approximation to this objective. Apart from that, Barlow Twins cannot yield ordered representations and handle graph representation learning. For the empirical evaluation, we have included comparisons to Barlow Twins in Tables 1-4. The results show statistically significant improvements over Barlow Twins across several tasks, batch sizes, projector dimensions, and datasets. We believe these gains are meaningful and kindly ask the reviewer to reconsider their evaluation of these results.

Difference between SCL and Neural Eigenmap

We clarify that SCL does not directly compute the top-k eigenfunctions of a kernel. Rather, SCL recovers the subspace spanned by the eigenfunctions, without determining their specific order or ranking. As explained in HaoChen et al. (2021) in the paragraph before Lemma 3.1, SCL learns the eigenfunctions up to a right transformation. Consequently, the representation learned by SCL does not exhibit an ordered structure as ours, which results in the performance gap in Fig 2.

Using 3D shapes, e.g. the Standford Bunny, for showing the approximation of the eigenfunctions

Thanks for the advice. Although the primary focus of this paper is not the approximation accuracy of NeuralEF, we will try offering visualizations of the learned eigenfunctions for 3d shapes in the final version. In the meantime, we kindly direct the reviewer to the original NeuralEF paper (Deng et al, 2022) for approximation results on 3d shapes like swiss rolls.