T-REGS: Minimum Spanning Tree Regularization for Self-Supervised Learning

We thank you for your feedback. Please find below our answers:

W1. “The theoretical analysis is based on strict sphere constraints, which cannot be completely ensured in reality. Since regularization is only able to gradually enforce the sphere constraints, a relaxed assumption is more suitable for the setting.”

Thank you for the nice and valuable suggestion! We have extended Proposition 4.1 to the case where the points lie inside the unit ball, not only on the sphere, with the same conclusion: the maximum length of the $\mathrm{MST}$ is achieved by points sitting on the sphere, at the vertices of a regular simplex. This new statement will replace Proposition 4.1 in the final version of the paper. Its proof is actually simpler than the previous one, relying on a standard result in convex geometry—see e.g. Chapter 14, Theorem 14.6, in [1]---to which we add a technical lemma relating the total sum of pairwise distances to the total length of the MST.

The new Proposition 4.1 explains the behavior of T-REG when used with a soft sphere constraint: first, minimizing the term $\mathcal L_E$ in Equation (7) expands the point cloud until the sphere constraint $\mathcal L_S$ becomes the dominating term; at that stage, the points stop expanding and start spreading themselves out uniformly along the sphere of directions. The amount of expansion before spreading is prescribed by the strength of the sphere constraint term versus the $\mathrm{MST}$ term in the loss, which is driven by the ratio between their respective mixing parameters $\lambda$ and $\gamma$. Note that $\mathcal L_E$ grows linearly (or sublinearly for $\alpha<1$) with the scale, whereas $\mathcal L_S$ grows quadratically, which explains why $\mathcal L_S$ eventually becomes the dominating term.

W2 & Q1 “There is still a gap between the performance of T-REGS and baselines, making its effectiveness not convincing enough. More trials regarding the proper choices of the parameters are needed.” / “Are there any results of combining T-REGS with existing SSL, both contrastive and non-contrastive? This would further demonstrate its effectiveness”

Given the feedback received, we conducted further experiments to combine T-REG with a couple of other regularizers: BYOL and Barlow Twins, on CIFAR-10/100. We provide the results in Table A.2, which highlight the benefits of such combinations (best performances in bold):

Table A.2: CIFAR-10/100 Evaluation

Adding our regularizer improves the results both for BYOL and Barlow Twins, highlighting its superiority.

W3. “Table 2: higher values of beta/gamma and gamma/lambda lead to better performance. Intuitively, this would mean that L_E and L_S are not that necessary. To justify the regularization terms, you should first show that no term has poorer performance than your method, and that extremely large and also harms the accuracy.”

Following your recommendation, we propose to complete Table 2 as shown below:

Table B.1: Impact of coefficients. We report Top-1 accuracy on ImageNet-1k.

The new version of the table (Table B.1) more clearly highlights that no regularization at all leads to collapse (line 1), that large values of regularization also harm classification accuracy (lines 9-10: 51.8% vs best performance of 66.1%), and thus that both $\mathcal L_E$ and $\mathcal L_S$ are essential for optimal performance.

Q2. “Can T-REGS be generalized to supervised learning?”

The regularization, T-REG, can indeed be applied to other settings beyond the one considered in the paper, when uniformity properties are helpful and/or where dimensional collapse can be an issue. While we do not have an example in supervised learning per se, we have other examples of applications in unsupervised or self-supervised settings.

For instance, Diffusion Models have been shown to suffer from dimensional collapse. Recently, Wang et al. [2] proposed to regularize a generative model's internal representations by encouraging their dispersion in the hidden space.

In multi-modal applications, several papers [3,4] have highlighted that uniformity properties can be highly beneficial to bridge the modality gap [5] in multi-modal models. The modality gap is a phenomenon observed in pre-trained multimodal models, such as CLIP [6], which provide transferable embeddings mainly by aligning images and text in the same embedding space. Better uniformity allows for the embeddings from the different modalities to more evenly distribute on the hypersphere in the embedding space. To validate the generalizability, we assess our method on image-text retrieval, a representative vision-language task, following the same experimental setting as Oh et al. (Section 5.2) [3] on Flickr30k and MS-COCO. More precisely, we apply T-REG as a regularizer to finetune CLIP: ($\mathcal L = \mathcal L_{\text{CLIP}} + \mathcal L_{\text{T-REG}}$), where $\mathcal L_{\text{CLIP}}$ is the standard CLIP loss [6] and $\mathcal L_{\text{T-REG}}$ refers to Eq. (7). We report the top 1/5 Recall (R@1, R@5) in Table B.2 below with Image-to-text (i → t) and text-to-image (t → i) retrieval results. Except for the last line, the table is inherited from Oh et al. [3] (bold indicates the best performance). Complete results and implementation details will be included in the supplementary material of the paper.

Table B.2: Image-to-text (i → t) and text-to-image (t → i) retrieval results

Q3. “What is the extra cost of MST computation?”

We use Kruskal’s algorithm, which runs in $\mathcal{O}(B^2 (D + \log B))$ time in the worst case, where $B$ is the number of samples in the batch and $D$ is the dimension, with a resulting computational overhead that is similar to well-established methods such as VICReg or SimCLR; and faster wall-clock time per step. We report the complexities in the table below, along with wall-clock time per step, following a suggestion by Reviewer MGro.

Table A.1: Compute comparison. The ranges are reported from Garrido et al [4] for SimCLR and Bardes et al [3] for VICReg.

---

References:

[1] T. M. Apostol and M. A. Mnatsakanian: “New Horizons in Geometry”, Mathematical Association of America, Dolciani Mathematical Expositions #47, 2012, ISBN 978-0-88385-354-2.

[2] R. Wang, K. He. Diffuse and Disperse: Image Generation with Representation Regularization. Arxiv 2025

[3] C. Oh, J. So, H. Byun, Y. Lim, M. Shin, J-.J. Jeon, K. Song. Geodesic Multi-Modal Mixup for Robust Fine-Tuning. NeurIPS 2023

[4] J. Kim, S. Hwang. Enhanced OoD Detection through Cross-Modal Alignment of Multi-Modal Representations. In CVPR, 2025

[5] V. W. Liang, Y. Zhang, Y. Kwon, S. Yeung, and J. Y. Zou. Mind the gap: Understanding the modality gap in multi-modal contrastive representation learning. NeurIPS, 2022

[6] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In Proc. ICML, 2021

[7] A. C. Li, A. A. Efros, and D. Pathak. Understanding collapse in non-contrastive siamese representation learning. In ECCV, 2022

[8] K. Lee, Y. Zhu, K. Sohn, C.-L. Li, J. Shin, and H. Lee. i-mix: A domain-agnostic strategy for contrastive representation learning. In ICLR, 2021.

[9] Z. Shen, Z. Liu, Z. Liu, M. Savvides, T. Darrell, and E. Xing. Un-mix: Rethinking image mixtures for unsupervised visual representation learning. In AAAI, 2022

[10] Q. Garrido, Y. Chen, A. Bardes, L. Najman, and Y. Lecun. On the duality between contrastive and non-contrastive self-supervised learning. In ICLR, 2023

[11] A. Bardes, J. Ponce, and Y. Lecun. Vicreg: Variance-invariance-covariance regularization for self-supervised learning. In ICLR, 2022.

We thank you for your feedback and positive comments. Please find below our answers to the weaknesses you raised:

W1. “1.Empirical Performance: While the results are competitive, T-REGS does not consistently set a new state-of-the-art on the reported benchmarks. For example, on CIFAR-100, it lags behind several W2-regularized methods, and on ImageNet-1k, it is slightly outperformed by methods like INTL and Zero-CL. While SOTA performance is not a strict requirement for a conceptually novel paper, a stronger showing would make the practical case for adoption more compelling. The authors acknowledge this limitation in the conclusion [305-306].“

In order to strengthen the empirical validation of our regularization method, we followed suggestions by Reviewers Hgsb and hUvh and conducted further experiments to integrate T-REG as an auxiliary regularizer to BYOL and Barlow Twins on CIFAR-10/100. The results are reported in Table A.2 below, which highlights the benefits of such combinations (best performances are shown in bold):

Table A.2: CIFAR-10/100 Evaluation

We observe that adding our regularizer improves the results both for BYOL and Barlow Twins, highlighting its superiority.

W2 “2.Computational Overhead of MST: The paper states that fast GPU-based MST implementations exist, but it does not provide a practical analysis of the computational overhead. For a batch of size B, constructing the complete graph has a complexity of O(B²), and standard MST algorithms like Kruskal's or Prim's are at least O(B²). This could become a bottleneck for very large batch sizes, which are common in SSL. A comparison of wall-clock training time per step against baselines like VICReg (which computes a covariance matrix, O(BxD²)) or BarlowTwins would be a valuable addition to assess the practical scalability.”

Indeed, the worst-case asymptotic complexity is significant, but not larger than for other methods. Meanwhile, the hidden constants in the complexity bounds are small, which has an impact in practice since the batch sizes are relatively small (only a few thousand points). Furthermore, while Kruskal’s algorithm per se is sequential, its pre-processing (namely, computing the distance matrix and sorting its entries) is the most costly step. In practice, one can use a parallelized GPU implementation for it. This is what the torchPH library (used in our implementation) does.

To confirm this insight, we evaluated the computational overhead of T-REGS, VICReg (a covariance-based method), and SimCLR (a contrastive method). We report the results in the table below, where $B$ represents the batch size and $D$ represents the output dimension. The wall-clock comparisons are averaged over 500 steps for each method.

Table A.1: Compute comparison. The ranges are reported from Garrido et al [4] for SimCLR and Bardes et al [3] for VICReg.

We observe a resulting computational overhead that is similar to well-established methods such as VICReg or SimCLR; and faster wall-clock time per step.

---

References:

[1] A. Bardes, J. Ponce, and Y. Lecun. Vicreg: Variance-invariance-covariance regularization for self-supervised learning. In ICLR, 2022.

[2] Q. Garrido, Y. Chen, A. Bardes, L. Najman, and Y. Lecun. On the duality between contrastive and non-contrastive self-supervised learning. In ICLR, 2023

We thank you for your feedback. Please find below our answers to the questions and weaknesses you raised, and let us know how we can help further clarify the method for you:

W1: “The authors did not provide a relatively clear explanation of how the Minimum Spanning Tree (MST) is combined with Self-Supervised learning. The picture provided by the authors in Figure 1 is common knowledge. Instead of providing a common-sense picture, the authors should have provided a picture in the introduction to help readers intuitively understand how the Minimum Spanning Tree (MST) is combined with Self-Supervised learning.”

We agree that Figure 1 in the introduction is standard, and that a more targeted illustration would help clarify how the MST is integrated into the SSL framework. We will modify the figure to better demonstrate how the minimum spanning tree is constructed from embeddings and how it fits into the pipeline. For this, we will take inspiration from the visuals of Figure 2, which will be moved from Page 7 to the introduction to provide a better initial intuition for the reader.

We will also add a few clarifying sentences to the introduction, in the same spirit as our answer to Q1 below.

Q1: “In lines 66-70, the author introduced the concept of alpha-minimum spanning tree, but the author didn’t combine it well with the context of self-supervised learning. I’m very curious about how to construct a connected acyclic graph given the embeddings of self-supervised learning.”

We kindly refer you to Equation (10) in Section 5, which explains this point.

Our regularization is applied separately on the embeddings of the two transformations of the batch obtained from the two network branches. Specifically, denoting by $Z=(z_1, \dots ,z_n)$ and $Z' = (z_1', \dots, z_n')$ these two embeddings in $d$ dimensions, two minimum spanning trees $\mathrm{MST}(Z)$ and $\mathrm{MST}(Z’)$ are constructed on these two sets of points. The MST of $Z$ treats each point $z_i$ as a vertex, and connects the vertices by edges to form the connected acyclic graph of minimum total length, where the length of an edge corresponds to the Euclidean distance between its vertices. Likewise for the MST of $Z'$. The regularization terms in Equation (10) tend to maximize the total lengths of these two minimum spanning trees while not expanding the point clouds too much, which provably increases their entropy and dimensionality (please refer to our answer to W2 and Q2 below for the explanation).

We intend to add one or two sentences in the same spirit as above to the introduction in order to clarify this point.

W2 & Q2 : “I didn’t understand how the Minimum Spanning Tree solves the problems of avoiding dimensional collapse and enhancing uniformity in self-supervised learning.”/ “I want to know how the Minimum Spanning Tree solves the problems of avoiding dimensional collapse and enhancing uniformity in self-supervised learning.”

The behavior of T-REG is explained by Proposition 4.1, which we have now extended to allow the points to lie inside the ball and not just on the sphere. First, minimizing the term $\mathcal{L}_E$ in Equation (7) expands the point cloud until the sphere constraint $\mathcal{L}_S$ becomes the dominating term; at that stage, the points stop expanding and start spreading themselves out uniformly along the sphere of directions, which prevents dimensional collapse and maximizes uniformity. The amount of expansion before spreading is prescribed by the strength of the sphere constraint term versus the $\mathrm{MST}$ term in the loss, which is driven by the ratio between their respective mixing parameters $\lambda$ and $\gamma$.

For further visual intuition, we kindly refer you to the synthetic experiments in dimension 3 (Figures 2.a and 2.b, which will be moved to the introduction). We generate a highly concentrated point cloud of 256 points (orange points in Figures 2a and 2b). After optimization, T-REG successfully transforms this initially concentrated point cloud into a uniformly distributed point cloud on the sphere (blue points in Figure 2.a), while occupying the whole dimensional space—demonstrating no dimensional collapse. Without the sphere constraint, the optimization fails to converge as points diverge to infinity (Figure 2.b).

While I appreciate the authors’ detailed responses, the critical issue concerning the integration of Minimum Spanning Trees (MST) with Self-Supervised Learning (SSL) remains inadequately addressed, both in the original manuscript and the rebuttal. I concur with Reviewer Hgsb regarding the lack of clarity.

Major Concern (W1 & Q1: MST-SSL Integration): The fundamental mechanism of how MST is applied within the SSL framework is still unclear. The authors’ responses to my concerns (W1, W2, Q1), particularly W1 and Q1 regarding the core concept and rationale of MST usage, did not resolve the ambiguity. This lack of intuitive understanding makes it difficult for the SSL community to grasp the novelty and contribution of using MST in this context.

It is essential that the revised manuscript provides a significantly clearer explanation. I strongly suggest supplementing the textual description with an intuitive visual illustration (e.g., diagram, flowchart) explicitly showing how the MST computation and structure are utilized within the SSL pipeline. Clarity on this point is paramount for acceptance.

Pending substantial improvement on this central issue, I cannot currently recommend acceptance. I maintain my score.

We thank you for your feedback. Please find below our answers to the questions and weaknesses you raised:

W1: “The writing of the paper can be greatly improved. E.g., proofs can be moved to the appendix. In the main paper, instead of showing proofs, explain why these theoretical results are important for its application here.”

We chose to include the proofs of the two pivotal results (Proposition 4.1, Proposition 4.3) in the main text because these results are central to the paper and explain the behavior of T-REG, with relatively short and elementary proofs. The final version of the paper will provide intuitive explanations of the role of the MST for preventing dimensional collapse and promoting well-distributed embeddings, based on Propositions 4.1 and 4.3—see our answer to Weakness 2 of Reviewer zFs8 for the details. This will add a few sentences to the introduction and Section 4, which will fit together with the new experimental results within the extra allocated page.

W2. “Comparing to other existing regularization techniques, T-REGS doesn't seem to result in better performance. It's unclear what's the benefit of adopting T-REGS is compared to other methods.

The main benefits of T-REGS stem from its conceptual and implementation simplicity, strong theoretical grounding, and robustness to hyperparameter variations (i.e., batch size and embedding dimension). Indeed, T-REGS offers a novel, conceptually simple solution based on well-understood graph theory and statistical dimension estimation. It follows the principled line of thought of Fang et al. [1] while proposing a more straightforward and unexplored approach: maximizing the length of the minimum spanning tree, which we found to satisfy the same properties naturally but with significantly lower computational overhead in practice and improved numerical stability. As a result, T-REGS demonstrates better scalability on larger datasets such as ImageNet-100 and ImageNet-1k.

Compared with other existing approaches, T-REGS addresses the following well-known limitations:

(i) Theoretical grounding: Asymmetric methods lack theoretical justification to explain how the architecture helps prevent collapse [2]. Covariance-based methods only leverage the second moment of the data distribution and are thus blind to the concentration points of the density, which can prevent convergence to the uniform distribution (Figure 5 in the Supp. Mat.),

(ii) Computational limitations: contrastive methods are sensitive to the number of negative samples and require large batch sizes [3, 4]. Covariance-based approaches often rely on both large batch sizes and large output dimensions [2]. By contrast, T-REGS proves to be robust to both batch size and embedding dimensions (Appendices D.1 and D.2).

To strengthen the empirical validation of our method, we have conducted complementary experiments based on your suggestions. We propose a revised version of Table 1 that highlights the new results (please refer to our answer to Q2).

Q1: “Could you elaborate how exactly the length of minimum spanning tree is computed? Which algorithm is used for finding the minimum spanning tree found given a batch of features and what's the complexity of the algorithm? Could you compare the complexity with other regularization techniques?”

In practice, we use Kruskal’s algorithm, which runs in $\mathcal{O}(B^2 (D + \log B))$ time in the worst case, where $B$ is the number of samples in the batch and $D$ is the dimension. This choice was motivated by the simplicity of its implementation, with a resulting computational overhead that is similar to well-established methods such as VICReg or SimCLR; and faster wall-clock time per step. We report the complexities in Table A.1 below, along with wall-clock time per step, following a suggestion by Reviewer MGro.

Table A.1: Compute comparison. The ranges are reported from Garrido et al [4] for SimCLR and Bardes et al [3] for VICReg.

Q2: “It's unclear to me why T-REGS only uses MSE as the representation learning loss. Why would it not be able to act as a generic regularization for all SSL methods?”

T-REG is indeed designed as a generic regularization technique and, in principle, can be combined with a wide range of self-supervised learning (SSL) objectives—not just as a standalone regularization with an MSE. In the submitted version of the paper, we decided to emphasize the behavior of T-REGS (T-REG combined with MSE invariance) as a standalone method. Given the feedback received, we have run further experiments combining T-REG with a couple of other regularizers: BYOL and Barlow Twins, on CIFAR-10/100. We provide the corresponding results below in Table A.2, which highlight the benefits of such combinations (best performances are shown in bold):

Table A.2: CIFAR-10/100 Evaluation

We observe that adding our regularizer improves the results both for BYOL and Barlow Twins, highlighting its superiority.

Q3: “Do you have any ablations for the value of alpha?”

We tested $\alpha=1$ and $\alpha=0.5$ and found that $\alpha=1$ works best; however, we have not performed an exhaustive hyperparameter sweep. Table A.3 reports the results:

Table A.3: Ablation of $\alpha$. Results on CIFAR-10/CIFAR-100.

Q4: “In many popular SSL framework, the features are normalized to unit sphere. It has also been shown such normalization is beneficial. Why does the paper propose to regularize it via the soft sphere-constraint instead of a hard constraint? Could you provide empirical evidence suggesting this way better?”

Indeed, some popular SSL frameworks such as SimCLR and BYOL employ explicit normalization of features to the unit sphere, enforcing a hard constraint. Others do not enforce such constraints explicitly, and instead rely on soft mechanisms—such as VICReg, which incorporates variance and covariance regularization terms in its loss, implicitly constraining the distribution (and to some extent, the norms) of the embeddings.

We chose a soft sphere constraint for the following reasons:

• A hard normalization has been shown to ignore the importance of the embedding norm for gradient computation, whereas a soft constraint enables better embedding optimization [5,6]
• Furthermore, during our initial experiment, we found that relaxing the sphere constraint from a hard one to a soft one provides more leeway for optimization and leads to improved results. Table A.4 shows the results with/without normalization.

Table A.4: Results with and without normalization

---

References:

[1] X. Fang, J. Li, Q. Sun, and B. Wang. Rethinking the uniformity metric in self-supervised learning. In ICLR, 2024

[2] X. Weng, Y. Ni, T. Song, J. Luo, R. M. Anwer, S. Khan, F. Khan, and L. Huang. Modulate your spectrum in self-supervised learning. In ICLR, 2024

[3] Q. Garrido, Y. Chen, A. Bardes, L. Najman, and Y. Lecun. On the duality between contrastive and non-contrastive self-supervised learning. In ICLR, 2023

[4] A. Bardes, J. Ponce, and Y. Lecun. Vicreg: Variance-invariance-covariance regularization for self-supervised learning. In ICLR, 2022.

[5] D. Zhang, Y. Li1, Z. Zhang. Deep Metric Learning with Spherical Embedding. NeurIPS 2020

[6] J. Zhang, H. Zhang, R. Vasudevan, M. Johnson-Roberson. Hyperspherical Embedding for Point Cloud Completion. In CVPR 2023.