🤖 machine learning

On the Importance of Embedding Norms in Self-Supervised Learning

This paper resolves the apparent contradiction regarding the role of embedding norms in self-supervised learning by demonstrating through theoretical and experimental analysis that these norms, despite the prevalence of cosine similarity, critically govern convergence rates and encode network confidence, where smaller norms indicate unexpected samples.

Original authors: Andrew Draganov, Sharvaree Vadgama, Sebastian Damrich, Jan Niklas Böhm, Lucas Maes, Dmitry Kobak, Erik Bekkers

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Andrew Draganov, Sharvaree Vadgama, Sebastian Damrich, Jan Niklas Böhm, Lucas Maes, Dmitry Kobak, Erik Bekkers

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Spiky Ball" Problem

Imagine you are teaching a computer to recognize pictures (like cats and dogs) without showing it any labels. This is called Self-Supervised Learning (SSL).

To do this, the computer turns every picture into a mathematical point in a giant, multi-dimensional space. To make the math work, the computer usually forces all these points to sit on the surface of a perfect, smooth ball (a hypersphere). It does this by measuring how similar two points are based on their direction (like checking if two arrows are pointing the same way), ignoring how long the arrows are.

The Paper's Discovery:
The authors found that even though the computer is supposed to ignore the length of these arrows (the "embedding norm"), the length actually matters a lot. In fact, the training process naturally makes the arrows for common, easy-to-recognize pictures grow very long, while the arrows for rare or confusing pictures stay short.

This turns the computer's "perfect smooth ball" into a spiky ball. The spikes are the long arrows for familiar data, and the flat areas are the short arrows for unfamiliar data.

The Two Main Rules of the Paper

The paper explains two main things about these "arrow lengths" (norms):

1. The "Heavy Backpack" Effect (Convergence Speed)

The Analogy: Imagine you are trying to walk toward a destination (learning the right answer). If you are carrying a heavy backpack (a long arrow/norm), you move much slower. If you are light (a short arrow), you can sprint.

What the Paper Says:
The math shows that the longer the arrow gets, the slower the computer learns. Specifically, the learning speed drops by the square of the arrow's length.

The Catch-22: To learn, the computer needs to push the arrows together. But the act of pushing them together naturally makes the arrows grow longer.
The Result: The computer gets stuck in a loop where it tries to learn, but its own learning makes the "backpack" heavier, slowing it down.

2. The "Confidence Meter" (Network Confidence)

The Analogy: Think of the arrow length as a volume knob for confidence.

Loud (Long Arrow): The computer is very sure about this picture. It sees this type of cat often, so it has a strong, long signal for it.
Quiet (Short Arrow): The computer is unsure. This might be a weird angle of a cat, or a picture of a dog that looks like a cat. The signal is weak and short.

What the Paper Says:
The length of the arrow naturally encodes how confident the model is.

Common Data: If a picture looks like the training data, the arrow gets long (High Confidence).
Strange Data: If a picture is totally new or weird (Out-of-Distribution), the arrow stays short (Low Confidence).
Imbalanced Data: If the computer sees "Cats" 1,000 times and "Dogs" only 10 times, the "Cat" arrows will grow huge, and the "Dog" arrows will stay tiny. This makes the computer biased and unstable.

The Solutions: How to Fix the Spiky Ball

The authors tested three ways to stop the arrows from growing uncontrollably and to fix the learning speed.

1. The "Weight Decay" (The Tether)

What it is: A standard tool in machine learning that gently pulls the weights back toward zero.
The Paper's Finding: It helps keep the arrows from getting too long, but it's a slow, gradual fix. If you pull too hard, the computer forgets everything (the ball collapses). If you don't pull enough, the arrows get too long and learning slows down.

2. The "Cut-Initialization" (The Starting Line)

What it is: Before the training even starts, the authors take all the computer's internal numbers and divide them by a constant (like 3 or 9).
The Analogy: Imagine starting a race with everyone wearing heavy boots. Instead, you start them barefoot.
The Paper's Finding: This is a huge win. By starting with short arrows, the computer learns much faster. It prevents the "heavy backpack" problem from the very first step. It works especially well for models that don't use "negative" examples (non-contrastive models like SimSiam).

3. The "GradScale" (The Volume Knob)

What it is: A special layer that changes how the computer learns. It looks at the length of the arrow and adjusts the learning step.
The Analogy: If the arrow is long (heavy backpack), the computer takes a tiny step. If the arrow is short, it takes a big step.
The Paper's Finding: This cancels out the "heavy backpack" effect entirely. It makes the learning speed consistent regardless of how long the arrow is. However, the paper notes this can be tricky to tune and sometimes causes the computer to get confused if the data is too messy.

Summary of Results

The Problem: SSL models naturally create "spiky" representations where common data has long arrows and rare data has short ones. This slows down learning and creates bias.
The Good News: The length of the arrow is actually a useful signal! It tells you how confident the model is.
The Fix: You can fix the speed and stability issues by:
1. Starting with smaller numbers (Cut-Initialization).
2. Using a special layer to adjust learning steps based on arrow length (GradScale).
3. Carefully tuning how much you pull the weights back (Weight Decay).

The paper concludes that we shouldn't just look at the direction of the data points; we must also manage their length to make self-supervised learning faster and more reliable.

Technical Summary: On the Importance of Embedding Norms in Self-Supervised Learning

Problem Statement

Self-supervised learning (SSL) has become a dominant paradigm for learning data representations without supervised signals, particularly through methods like SimCLR, MoCo, SimSiam, and BYOL. These methods typically rely on optimizing the cosine similarity between embedding vectors, effectively constraining the learning dynamics to a hypersphere. Consequently, theoretical analyses have largely focused on the distribution of points on this hypersphere, often assuming that embedding norms (the magnitude of the vectors before normalization) are irrelevant or merely a byproduct of $L_2$ normalization.

However, empirical observations suggest that embedding norms contain meaningful information related to network convergence and confidence. Previous works noted that norms inversely scale gradients under cosine similarity loss and that they correlate with model certainty, yet a systematic theoretical and empirical understanding of this relationship was lacking. This paper addresses the apparent contradiction: while SSL optimizes for direction (cosine similarity), the magnitude of the embedding vectors plays a critical, often detrimental, role in training dynamics.

Methodology

The authors employ a combination of theoretical analysis, controlled simulations, and extensive empirical experiments to investigate the role of embedding norms in SSL.

Theoretical Analysis:
- The paper derives the gradients of the attractive term ( $L^A_{ij}$ ) and the InfoNCE loss with respect to an embedding vector $z_i$ .
- It proves that gradients for cosine-similarity-based losses are orthogonal to the embedding vector itself.
- Using these properties, the authors demonstrate that applying gradient descent steps to maximize cosine similarity inevitably increases the magnitude (norm) of the embedding vectors.
- Theoretical bounds are established showing that the change in cosine similarity per gradient step is inversely proportional to the square of the embedding norm ( $\rho^2$ ). This implies that as norms grow, the effective learning rate decreases quadratically, slowing convergence.
Simulations:
- Idealized Convergence: The authors simulate gradient descent on pairs of points on a hypersphere to verify that larger initial embedding norms lead to significantly slower convergence, confirming the quadratic slowdown predicted by theory.
- Latent Class Dynamics: Using a synthetic dataset generated from latent classes, the study simulates SSL training to observe how embedding norms evolve. They find that embeddings in dense regions of the latent space (frequently observed samples) grow larger norms than those in sparse regions.
Empirical Experiments:
- Confidence Correlation: Experiments on CIFAR-10, CIFAR-100, and ImageNet-100 demonstrate that embedding norms correlate with model confidence. Samples with higher norms are more likely to be in-distribution, correctly classified by a kNN classifier, and associated with high human labeler agreement (low entropy). Conversely, out-of-distribution (OOD) samples exhibit significantly smaller norms.
- Intervention Strategies: The paper evaluates three mechanisms to mitigate the "embedding-norm effect" (the tendency for norms to grow and slow training):
  - Weight Decay: Standard regularization that penalizes large weights, indirectly controlling embedding growth.
  - Cut-Initialization: A novel initialization strategy where network weights are divided by a constant $c$ at the start of training to ensure small initial embedding norms, which are then maintained by weight decay.
  - GradScale Layer: A custom PyTorch autograd function that rescales the gradient by the embedding norm to the power of $p$ . Setting $p=1$ cancels the norm's inverse scaling effect on the gradient, effectively normalizing the learning rate per sample.

Key Contributions and Results

1. The Embedding-Norm Effect and Convergence

The paper establishes that optimizing cosine similarity creates a "catch-22": small embedding norms are required for fast convergence (to avoid vanishing gradients), yet the optimization process itself drives these norms to grow.

Result: Theoretical bounds and simulations confirm that embedding norms impose a quadratic slowdown on SSL convergence.
Result: Weight decay can mitigate this, but excessive decay leads to representation collapse, while insufficient decay allows norms to grow unchecked, slowing training.

2. Embedding Norms as Confidence Metrics

The authors demonstrate that embedding norms naturally encode a model's confidence in the learned representation.

Result: High embedding norms correspond to frequent, prototypical samples and high classification accuracy.
Result: Low embedding norms correspond to rare samples, out-of-distribution data, and samples with high human labeler disagreement (high entropy).
Result: This phenomenon holds across both contrastive (SimCLR) and non-contrastive (SimSiam) methods, suggesting it is an intrinsic property of cosine-similarity-based SSL.

3. Mitigation Strategies and Performance Gains

The paper proposes and validates methods to manage embedding norms, showing significant improvements in training efficiency and final accuracy.

Cut-Initialization: Pairing cut-initialization with standard weight decay accelerates convergence and improves final kNN accuracy. For example, on CIFAR-100, SimCLR with cut-initialization ( $c=3$ ) improved accuracy by ~1-2% over the default, while SimSiam saw gains of ~10% on ImageNet-100.
GradScale: This method, which explicitly cancels the norm-dependent gradient scaling, provides further improvements, particularly in class-imbalanced settings where variance in embedding norms can destabilize training.
Imbalanced Data: The embedding-norm effect is particularly detrimental in class-imbalanced datasets. The proposed interventions (Cut-Initialization and GradScale) yield double-digit accuracy improvements in these scenarios (e.g., improving SimSiam on imbalanced CIFAR-100 from 21.6% to 31.5%).
Robustness: These benefits persist even when using adaptive optimizers like Adam.

Significance and Claims

The paper claims to resolve a gap in the literature by providing the first thorough analysis of how embedding norms interact with SSL training dynamics. The authors argue that embedding norms are not merely a byproduct of normalization but are integral to understanding and optimizing network behavior.

Theoretical Insight: The work clarifies the "embedding-norm effect," showing that it is an inevitable consequence of cosine-similarity-based optimization that slows convergence quadratically.
Practical Utility: The findings suggest that managing embedding norms is essential for effective SSL training, especially in non-contrastive settings and imbalanced data scenarios. The proposed "Cut-Initialization" and "GradScale" mechanisms offer simple, principled ways to improve training stability and performance.
Interpretability: The paper posits that embedding norms serve as a natural, intrinsic measure of model confidence and sample novelty, offering a new lens for analyzing SSL representations without requiring additional supervised signals.

The authors conclude that while the embedding-norm effect presents challenges, it also offers opportunities, such as using norms to guide training on unseen inputs or to identify prototypical examples. They note that non-contrastive architectures appear more sensitive to these effects, raising questions for future work regarding architectural design and the potential benefits of the effect in specific transfer-learning contexts.