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

AI summary

Overview

• Research area: Self-supervised learning (SSL), specifically joint-embedding SSL (JE-SSL) for vision, with the regularization criterion drawn from Topological Data Analysis and geometric measure theory.
• Technical level: Advanced. The paper combines a practical training loss with proofs about regular simplices, Riemannian manifolds, Rényi entropy, and MST growth rates.
• Scope: The paper proposes T-REG, a loss that maximizes the length of the Minimum Spanning Tree (MST) of a batch of embeddings while softly constraining those embeddings to a sphere, and extends it to T-REGS for SSL, with theory and experiments on synthetic data, CIFAR-10/100, ImageNet-100/1k, and CLIP-based image-text retrieval.
• Authors and affiliations: Julie Mordacq, David Loiseaux, Vicky Kalogeiton, Steve Oudot (Inria Saclay; LIX, CNRS, École Polytechnique, IP Paris). Posted to arXiv (2510.23484v2 [cs.LG], dated 06 Feb 2026), CC BY 4.0.

What This Paper Is About

Methods that learn representations from unlabeled images face two failure modes: dimensional collapse, where embeddings occupy only a low-dimensional subspace of the representation space, and poor uniformity, where embeddings bunch up in a few regions instead of spreading over the space. The paper's goal is a regularization term that provably avoids both at once, is cheap to compute, and can be bolted onto existing SSL methods.

Key Contributions

1. T-REG (Equation 6): a regularization technique combining an MST-length maximization loss (Equation 4) with a soft hypersphere constraint (Equation 5), where the MST of the batch is treated as a differentiable function of the point positions (gradient given in Equation 3).
2. Theory and empirical evidence that T-REG prevents dimensional collapse and enforces sample uniformity: a finite-sample result (Theorem 4.1) and an asymptotic result on compact Riemannian manifolds (Equation 8, Proposition 4.5, Corollary 4.6), plus synthetic experiments.
3. T-REGS (Equation 10): an extension of T-REG to joint-embedding SSL by applying the regularizer independently to each augmented view branch, usable either standalone alongside a Mean Squared Error view-invariance term or as an auxiliary loss added to an existing SSL objective.
4. Benchmark evaluation of T-REGS on CIFAR-10/100, ImageNet-100/1k, and CLIP image-text retrieval on Flickr30k and MS-COCO, plus a computational-cost analysis showing per-step wall-clock comparable to VICReg and SimCLR.

Main Findings

• Dimensional collapse is penalized: On 10,000 synthetic data points in dimension 1024 drawn from an isotropic Gaussian with a fraction η of coordinates zeroed out, the sensitivity of the T-REG loss to η was similar to that of the 𝒲₂ loss of Fang et al. [23], and the metrics −𝒲₂ and −ℒ_E decrease jointly as the collapse level η increases (Figure 4).
• Uniformity on synthetic point clouds: Starting from 256 points on a 1-d curve, T-REG spreads them uniformly over the sphere; using the MST-length term ℒ_E alone causes excessive dilation and fails to converge (Figure 3). In a 256-dimensional experiment with 256 points initialized as xᵢ = e₁ + εᵢ (εᵢ uniform in a ball of radius 0.001), the cosine-similarity distribution goes from a sharp peak near 1 to an almost-Dirac distribution slightly below 0, consistent with the vertices of a regular simplex (Figure 3, Theorem 4.1).
• Theory for small batches: Theorem 4.1 states that for n ≤ d+1 points confined to a closed Euclidean d-ball of radius r, the MST length is maximized when the points sit on the bounding sphere at the vertices of a regular (n−1)-simplex whose smallest circumscribing sphere is that ball, with edge length a = r·√(2(k+1)/k).
• Theory for large samples: For n → ∞, the rescaled MST length converges almost surely to C′∫f^((d−1)/d) dμ (Theorem 4.4, from [15, Corollary 5]), a quantity tied to the intrinsic Rényi (d−1)/d-entropy; Proposition 4.5 and Corollary 4.6 show the uniform density uniquely maximizes ∫f^p dμ for 0 < p < 1.
• CIFAR-10/100 (ResNet-18, 500 epochs, batch size 256, linear probing, Top-1): Barlow Twins improves from 91.2/68.2 to 91.8/68.5 with ℒ_T-REGS; BYOL improves from 89.5/63.7 to 90.4/65.7. Standalone ℒ_MSE + ℒ_T-REGS reaches 91.3/67.4, and the paper reports that T-REGS standalone lands within 0.1% of the best 𝒲₂-regularized approach on CIFAR-10.
• ImageNet-100 / ImageNet-1k (Top-1): BYOL + ℒ_T-REGS goes from 80.3/66.5 to 80.8/67.2; Barlow Twins + ℒ_T-REGS from 80.2/67.7 to 80.9/67.8; ℒ_MSE + ℒ_T-REGS reaches 80.3/68.8. For comparison, INTL is listed at 81.7/69.5.
• Image-text retrieval with CLIP: Fine-tuning with ℒ_CLIP + ℒ_T-REGS reports the best results in Table 3: Flickr30k R@1/R@5 of 83.2/96.0 for image-to-text and 80.8/96.4 for text-to-image; MS-COCO R@1/R@5 of 41.6/68.7 and 41.5/68.7.
• Hyperparameters matter: Table 4 reports β = 10, γ = 0.2, λ = 8e-4 (ratios β/γ = 50, γ/λ = 250) giving the best 50-epoch online-probe ImageNet-1k top-1 accuracy of 66.1; settings with β = 1 collapsed, and the paper notes that β ≥ γ ≥ λ was essential to prevent collapse.
• Compute cost is comparable: The MST-based loss costs 𝒪(B²(D·log B)) versus 𝒪(B²·D) for SimCLR and 𝒪(B·D²) for VICReg. Measured step times on ImageNet-1k across 4 Tesla H100 GPUs averaged over 500 steps (B = 512, D = 1024) were 0.22 ± 0.03 s for SimCLR, 0.23 ± 0.02 s for VICReg, and 0.20 ± 0.001 s for ℒ_MSE + ℒ_T-REGS.

Methodology in Plain English

Take a batch of embeddings produced by a network and build the Minimum Spanning Tree over them — the cheapest set of edges that connects all the points. Measure that tree's total length and push it upward. Because each point is pulled away from the points it is connected to in the tree, the batch expands. Left unchecked, the points would fly apart forever, so a second term penalizes any point whose distance from the unit sphere differs from 1. The two forces combine to spread the points evenly over the sphere instead of letting them pile up.

The MST used is unique under a genericity condition on the point cloud, and Kruskal's algorithm is used to compute it; the authors use the torchph library for GPU-parallel MST computation, noting that the distance-matrix computation and sorting dominate and parallelize well. In T-REGS, each of the two augmented view batches gets its own MST and its own regularization term, which makes the method architecture-agnostic and applicable to branches that differ in modality, as in the CLIP experiment. Training code is built on solo-learn, and the regularizer is combined with either a plain MSE view-invariance term or the original objective of BYOL, Barlow Twins, or CLIP using a mixing parameter β. Default standalone coefficients were β = 10, γ = 0.2, λ = 8e-4.

Why This Matters

Impact on research. The paper links a simple, cheap geometric regularizer to established results on MST length growth, Rényi entropy, and persistent-homology dimension, arguing that the same criterion that prevents dimensional collapse also drives the embedding distribution toward uniformity on arbitrary compact Riemannian manifolds. It is positioned as an alternative to Optimal Transport-based regularization (𝒲₂ from Fang et al. [23]), whose closed-form acceleration holds only on the sphere and requires square roots over SVD computations that may be numerically unstable. The authors state that T-REG satisfies the same four principled properties — instance permutation, instance cloning, feature cloning, and feature baby constraints — as that framework.

Real-world applications (these follow from the evaluated tasks rather than from deployed systems described in the paper):

• Large-scale image representation learning where labeled data is scarce or expensive, using SSL backbones pre-trained without annotations.
• Multimodal retrieval systems, where the reported gains are on image-to-text and text-to-image search over Flickr30k and MS-COCO.
• Transfer learning pipelines that reuse frozen pre-trained embeddings for downstream classification via linear probing.
• Training setups constrained by memory, since the method avoids the large negative-sample batches that contrastive approaches require.

Industry relevance. The method is presented as a drop-in auxiliary loss on top of existing SSL objectives such as BYOL and Barlow Twins, and the reported per-step wall-clock on ImageNet-1k is comparable to VICReg and SimCLR under the same hardware. That combination of drop-in integration and no meaningful step-time penalty is what makes it plausible for existing training pipelines. The paper does not report any deployed production system or inference-time cost measurements.

Future Directions

The paper's conclusion does not list explicit future work, so the following are open questions the work raises:

• Beyond the sphere. The theory is stated for compact Riemannian manifolds, but the practical implementations in this paper use a hypersphere with a soft constraint. Whether hard constraints or other manifolds change the behavior is not resolved.
• Larger-batch and scaling behavior. Theorem 4.1 covers the regime n ≤ d+1; ImageNet-1k experiments used ResNet-50 for 100 epochs, leaving open how T-REGS behaves at much longer schedules and larger backbones, which the paper does not report.
• Better MST computation. MST computation is built on Kruskal's algorithm with a sequential main loop; although preprocessing dominates and parallelizes, the authors present the 𝒪(B²(D·log B)) cost as the price of the method, leaving room for faster approximate or hierarchical constructions.
• Broader multimodal use. The CLIP experiment regularizes text and image branches independently. Whether the same term helps other paired-modality settings beyond retrieval, and whether it can close the modality gap rather than only improve uniformity, is not established in the reported results.

Target Audience

Researchers and practitioners in self-supervised representation learning who already know contrastive and non-contrastive methods such as SimCLR, MoCo, BYOL, Barlow Twins, VICReg, and Zero-CL, and who want a regularization criterion with theoretical backing. It is also suited to readers from Topological Data Analysis interested in an applied use of MST length and persistent-homology dimension, and to engineers looking for a drop-in auxiliary loss whose reported step time matches existing methods. The mathematical sections on regular simplices, Rényi entropy, and density optimization assume comfort with convex geometry and measure-theoretic arguments; readers without that background can still follow the loss definitions, the synthetic experiments, and the benchmark tables.

Authors’ abstract

Self-supervised learning (SSL) has emerged as a powerful paradigm for learning representations without labeled data, often by enforcing invariance to input transformations such as rotations or blurring. Recent studies have highlighted two pivotal properties for effective representations: (i) avoiding dimensional collapse-where the learned features occupy only a low-dimensional subspace, and (ii) enhancing uniformity of the induced distribution. In this work, we introduce T-REGS, a simple regularization framework for SSL based on the length of the Minimum Spanning Tree (MST) over the learned representation. We provide theoretical analysis demonstrating that T-REGS simultaneously mitigates dimensional collapse and promotes distribution uniformity on arbitrary compact Riemannian manifolds. Several experiments on synthetic data and on classical SSL benchmarks validate the effectiveness of our approach at enhancing representation quality.