JEPAMatch: Geometric Representation Shaping for Semi-Supervised Learning

1. Problem Statement

Semi-Supervised Learning (SSL) aims to leverage large amounts of unlabeled data to improve model performance when labeled data is scarce. While methods based on FixMatch (combining weak/strong augmentations with confidence-based pseudo-labeling) have achieved state-of-the-art results, they suffer from two critical bottlenecks:

1. Class Imbalance and Bias: Pseudo-label selection relies on fixed or adaptive confidence thresholds. This often leads to majority classes dominating the learning process, as they easily exceed the threshold, while minority classes receive sparse supervision. Incorrect pseudo-labels amplify this bias, leading to skewed models.
2. Slow Convergence: Early in training, model confidence is low, causing most unlabeled samples to be discarded (masked out). The model relies heavily on the small labeled set for many iterations, resulting in inefficient use of computational resources and requiring a large number of training steps (often >220k iterations) to converge.
3. Geometric Collapse: Standard SSL methods focus on discrete classification but often fail to structure the latent feature space effectively, leading to "dimensional collapse" where representations lose isotropic properties and fail to generalize well.

2. Methodology: JEPAMatch

The authors propose JEPAMatch, a paradigm shift from purely logical output thresholding to explicit geometric representation shaping. The method integrates the LeJEPA (Latent-Euclidean Joint-Embedding Predictive Architecture) framework with the adaptive thresholding of FlexMatch.

The learning process is decomposed into two synergistic levels:

A. Curriculum Level (Pseudo-Label Selection)

• Goal: Select reliable pseudo-labels while mitigating class imbalance.
• Mechanism: Utilizes FlexMatch, which employs class-wise dynamic confidence thresholds ($\tau_c$) rather than a single global threshold.
• Process:
  • Weakly augmented views ($x^w$) generate predictions.
  • Dynamic thresholds determine if a prediction is confident enough to be a pseudo-label.
  • Strongly augmented views ($x^s$) are trained using these pseudo-labels via cross-entropy loss.
  • This level ensures the classifier receives high-quality supervision without being biased toward majority classes.

B. Representation Level (Geometric Shaping)

• Goal: Structure the latent feature space to be robust, invariant, and isotropic, preventing dimensional collapse.
• Mechanism: Based on LeJEPA, which posits that high-quality representations should follow an isotropic Gaussian distribution.
• Architecture:
  • Multi-View Strategy: For each unlabeled image, the model generates:
    1. Weak View ($x^w$): Mild augmentation (global context).
    2. Strong View ($x^s$): Heavy augmentation (global context).
    3. Local Views ($x^{loc}$): $K$ smaller, local crops with mild augmentation.
  • Self-Supervised Prediction Loss ($L_{pred}$): The model predicts the representation of the weak view ($z^w$) using the strong view ($z^s$) and local views ($z^{loc}$). This forces the encoder to learn geometric consistency between global and local parts.
  • Adaptive Class-Wise SIGReg: To prevent dimensional collapse while maintaining class separability, the authors introduce a two-phase regularization strategy:
    1. Warmup Phase: Applies Global SIGReg (Sketched Isotropic Gaussian Regularization). It forces the distribution of local features to match a standard normal distribution $N(0, I)$. This ensures the network utilizes all dimensions and builds a non-collapsed space before reliable class clusters exist.
    2. Main Phase: Shifts to Adaptive Class-Wise SIGReg.
       • Class Means: Computes the mean $\mu_c$ for each class using labeled and confident pseudo-labeled data.
       • Centering: Features are centered relative to their class mean.
       • Variance Annealing: The target variance $\sigma$ is annealed from 1.0 down to 0.1, forcing features to form tight, isolated Gaussian clusters for each class.
       • Active Repulsion: A loss term ($L_{repulsion}$) penalizes the similarity between different class means to prevent them from collapsing into a single point.

C. Total Loss

The final objective combines supervised loss, unsupervised consistency loss, and the representation loss:
$$\mathcal{L}_{total} = \mathcal{L}_{sup} + \lambda_{unsup}\mathcal{L}_{unsup} + \lambda_{rep}\mathcal{L}_{rep}$$
Where $\mathcal{L}_{rep}$ includes the prediction loss, SIGReg, and repulsion loss.

3. Key Contributions

1. Paradigm Shift: Moves SSL from relying solely on confidence thresholds to explicitly shaping the geometry of the latent space using LeJEPA principles.
2. Two-Level Decoupling: Separates the discrete classification task (Curriculum Level) from the geometric organization of features (Representation Level), allowing each to be optimized independently yet synergistically.
3. Adaptive Class-Wise SIGReg: A novel regularization technique that transitions from a global isotropic constraint (to prevent collapse) to class-specific tight clusters (to improve discrimination), effectively solving the class imbalance issue inherent in pseudo-labeling.
4. Multi-View Geometry: Introduces a strategy using local crops alongside global views to enforce structural consistency in the latent space.

4. Experimental Results

The method was evaluated on CIFAR-100, STL-10, and Tiny-ImageNet.

• Performance: JEPAMatch consistently outperformed state-of-the-art baselines (FixMatch, FlexMatch, FreeMatch, SoftMatch, SimMatch, etc.).
  • On CIFAR-100 (400 labeled), it achieved 34.25% error rate (vs. 37.98% for FreeMatch).
  • On STL-10 (40 labeled), it achieved 13.44% error rate.
  • It showed significant robustness under extreme class imbalance (imbalance factors $\gamma=20, 50, 100$).
• Convergence Speed:
  • JEPAMatch converges 3 times faster than FlexMatch.
  • It reaches peak accuracy roughly 50,000 iterations earlier than FlexMatch.
  • It requires only 217 iterations (vs. the standard 220k+ for many SSL methods), drastically reducing computational costs.
• Pseudo-Label Quality: JEPAMatch maintains a significantly higher accuracy of assigned pseudo-labels throughout training (e.g., ~50% vs. ~35% for FlexMatch on CIFAR-100), indicating better noise filtering.

5. Significance

• Efficiency: By accelerating convergence, JEPAMatch makes semi-supervised learning more practical for resource-constrained environments.
• Robustness: The geometric shaping approach effectively mitigates the "confirmation bias" and class imbalance issues that plague traditional confidence-based SSL methods.
• Theoretical Grounding: By integrating LeJEPA, the paper bridges the gap between self-supervised representation learning (which focuses on structure) and semi-supervised classification (which focuses on labels), suggesting that structuring the latent space is as crucial as the labeling strategy itself.

In summary, JEPAMatch demonstrates that explicitly enforcing an isotropic, well-structured geometric representation space allows SSL models to learn faster, handle class imbalance better, and achieve higher accuracy with fewer labeled samples and training iterations.