Paper: Neural Eigenfunctions Are Structured Representation Learners (ICLR-2024)

Reviewer comment (Reproducibility):
I have tracked the training process over graph dataset arXiv. It seems to me that there is no special ways to take the augmented nodes. It seems to me it quite arbitrarily to construct X^+, although adjacency matrix is applied.

Author response:
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.

Reviewer comment (Reproducibility):
Not sure why the first term in (12) is squared in the code.

Author response:
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.

Reviewer comment (Writing):
Some of the choices made are not clearly explained. For example, the choice of enabling/disabling stop_grad is said to allow/not allow for ordered eigenfunctions/structured representations, but there is no discussion of this (e.g., what does stop_grad do).

Author response:
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.

Reviewer comment (Experiments):
The lack of explanation about stop_grad is also implied in the choice of elements with small indices vs. random elements being chosen when evaluating these two approaches.

Author response:
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.

Reviewer comment (Experiments):
The proposed method is similar in formulation to Barlow Twins, but the numerical comparison is focused on a few experiments.

Author response:
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.

Reviewer comment (Writing):
Just before eq. 3, define BN layers (as batch normalization layers?).

Author response:
Thanks. We have revised the paper accordingly.

Reviewer comment (Writing):
Section 4: "principle eigenfunctions" should be "principal eigenfunctions".

Author response:
Thanks. We have revised the paper accordingly.

Reviewer comment (Theory):
Can the optimization in (3) be connected to the original definition of the eigenfunction from (1)?

Author response:
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.

Reviewer comment (Theory):
Can the authors elaborate on how the ordered structure arises in NeuralEF? Particularly given that the orthogonality requirement for eigenvectors in an eigendecomposition is not present in the NeuralEF formulation (8) to obtain the set of "learned" eigenvectors.

Author response:
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).

Reviewer comment (Writing):
After eq. 10, should $\gamma$ be $\lambda$?

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

Reviewer comment (Evaluation):
Is there intuition as to why (9) works better than (10)?

Author response:
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.

Reviewer comment (Evaluation):
Can the authors elaborate on the scalability of the proposed approach vs. that of Johnson et al. (2022)?

Author response:
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.

Reviewer comment (Writing):
Section 6.2 first sentence refers to Section 6.2 - is that a typo?

Author response:
We apologize for the typo. We have revised the manuscript.

Reviewer comment (Reproducibility):
Maybe not a weakness, but in this kind of learning, how do we know which kernel to pre-define?

Author response:
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.

Reviewer comment (Experiments):
With respect to the linear probe experiments for unsupervised representation learning: Why have the authors reproduced Barlow Twins themselves when, if I am not mistaken, the top-1 accuracy for ImageNet with a ResNet-50 pretrained encoder is available (like it is for SCL, which the authors use)?

Author response:
*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.

Reviewer comment (Evaluation):
With respect to the linear probe experiments for unsupervised representation learning: Can you make any inferential comments on the link between batch size, number of epochs, projector dimensions?

Author response:
*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.

Reviewer comment (Experiments):
With respect to the linear probe experiments for unsupervised representation learning: For Barlow Twins, the authors in that paper report that the model has highest top-1 accuracy at a batch size of 1024, so why are the results selected for a batch size of 2048 in this paper?

Author response:
*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.

Reviewer comment (Theory):
Maybe something I am missing: while talking about the connection between PCA and the proposed method in the subsection on “Learning ordered representations”, I have a general doubt— isn’t PCA a special, linear case of spectral analysis on Euclidean space— and hence follows the observation on ordering and the principal components carrying the most information, or is there a deeper connection to talk about here?

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

Reviewer comment (Theory):
Something I didn’t understand: why does replacing $\hat{\psi_{X_B}}$ by $\psi_{X_B}$ not affect the optimal classifier? (when talking about the linear probe evaluation)

Author response:
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.

Reviewer comment (Experiments):
What are the benchmarks that ignore, and that which do not ignore feature importance? Are they based on the task of image retrieval?

Author response:
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.

Reviewer comment (Theory):
The authors remark that the proposed method has faster forward/backward passes than graph neural networks which have a computational complexity of $O(n^3)$-- what is the computational complexity of this method?

Author response:
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.

Reviewer comment (Novelty):
Regarding the novelty, it is unclear to me what is the difference between the proposed approaches compared to previous works, which makes hard to understand the actual contributions.

Author response:
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.

Reviewer comment (Novelty):
Regarding the graph-based problem (Sec. 4) it is clear that Eq. 12 is the proposed model, and due to $\hat{\psi}$ in the second term, the learned eigenfunctions are ordered. I think in Sec. 3 this is not entirely clear. As far as I understand, Eq. 8 which gives ordered representations is already proposed by Jonson et al. (2022)? While the proposed Eq. 9 does not give ordered representations but it is potentially scalable?

Author response:
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.

Reviewer comment (Writing):
I think that it is not clear from the text when the $\hat{\psi}$ in the second term of Eq. 8 should be used and when not.

Author response:
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.

Reviewer comment (Novelty):
I believe that the analysis of the comparison to Barlow Twins is rather limited. It seems that the approaches are very similar both in their formulation but also in the experiments. I believe the differences should be investigated in more detail.

Author response:
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.

Reviewer comment (Evaluation):
As far as I understand, both SCL and Neural Eigenmaps learn the top-k eigenfunctions of a kernel. Why there is a difference in performance e.g. Fig 2? Neural eigenmaps approximates better the true eigenfunctions compared to SCL?

Author response:
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.

Reviewer comment (Experiments):
There is an example in the appendix that shows the approximation of the eigenfunctions for the RBF kernel. Similar to the spectral methods, I think that is interesting to see a comparison using 3D shapes, e.g. the Standford Bunny.

Author response:
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.