Deep Learning · Representation Learning

Alignment and Uniformity as Separate Objectives

Contrastive losses can be decomposed into a term that pulls positive pairs together and a term that spreads all representations apart, and the two terms trade against each other in measurable ways.

Two properties are usually wanted from an embedding space: inputs that mean the same thing should have the same vector, and vectors for different inputs should not bunch together. In a contrastive loss these properties are coupled through the softmax, which makes it hard to say which one a given training run has achieved and which one is limiting performance. Writing the objective as two explicit terms separates them, and the separation turns out to be more than a change of notation.

The decomposition

Consider representations normalized to the unit sphere and a positive pair distribution. Define alignment as the expected squared distance between the representations of a positive pair:

L_align = E[ || f(x) - f(x+) ||^2 ]

and uniformity as the logarithm of the average pairwise Gaussian potential over all pairs:

L_uniform = log E[ exp( -2 || f(x) - f(y) ||^2 ) ]

The first term is small when positives map to the same point. The second is small when the representations are spread over the sphere, because the average of a decreasing exponential of the squared distance is smallest when the distances are large. The contrastive loss with a particular negative-sampling distribution is asymptotically equivalent to a weighted sum of these two, up to constants.

The equivalence is approximate and holds in the limit of many negatives. It is useful because it names the two mechanisms separately. A model can be well aligned and badly uniform, or the reverse, and the two failures look different in the space even though both raise the contrastive loss.

Alignment alone has a degenerate minimum

Minimizing L_align with no constraint drives every input to the same point. The minimum is zero and it is reached by a constant function. This is why the alignment term alone is not a training objective; it needs the uniformity term or an architectural asymmetry to prevent it.

Alignment is also the term that the augmentation policy controls most directly. If positives are mild augmentations of each other, alignment is easy and the encoder can achieve it by being insensitive to small changes. If positives are aggressive, alignment requires the encoder to be insensitive to large changes, which is a stronger claim about what the representation should contain. The augmentation choice described in Data Augmentation as an Invariance Specification is therefore a choice about how much alignment is being demanded.

Uniformity is about the whole sphere, not the batch

The uniformity term is an average over all pairs, but any finite batch estimates it from the pairs in the batch. With N items there are N(N-1)/2 distinct pairs, and a batch of 256 gives about 32,000 pairs, which is enough to estimate a smooth average over a sphere in a few hundred dimensions only roughly. The estimate is biased toward uniformity within the batch’s own region, because the batch contains no information about the rest of the space.

This is the mechanism behind the observation that contrastive training benefits from large batches. A larger batch gives a better estimate of the uniformity term and also covers more of the sphere, so the encoder is pushed to spread over regions it has actually seen. Methods that maintain a queue of past representations, described in Momentum Encoders and Stop-Gradients, extend the estimate beyond the current batch for the same reason.

Uniformity has a failure mode of its own: it can be satisfied by spreading points into directions that carry no information. A space that is uniform because the encoder outputs a fixed pseudo-random function of the input, uncorrelated with content, has low uniformity loss and no alignment. The two terms are only jointly meaningful.

Measuring the two terms separately

Both quantities can be computed on a held-out set without training, which makes them diagnostic rather than merely descriptive. Alignment is an average over positive pairs; uniformity is an average over all pairs, which for a large set is expensive but can be estimated from a random subset.

import torch
import torch.nn.functional as F

def align_uniform(z1, z2, t=2.0):
    z1 = F.normalize(z1, dim=1)
    z2 = F.normalize(z2, dim=1)
    align = (z1 - z2).pow(2).sum(1).mean()          # small = positives agree
    d2 = torch.cdist(z1, z1).pow(2)                  # squared euclidean
    mask = ~torch.eye(len(z1), dtype=torch.bool)
    uniform = (-t * d2[mask]).exp().mean().log()     # small = points spread
    return align.item(), uniform.item()

Reading the pair of numbers tells which term a training run has failed on. A low alignment and high uniformity indicates positives are not being matched, which points at the augmentation policy or at an optimization problem. A high alignment and high uniformity indicates the positives match but the space is crowded, which points at batch size, queue size, or a temperature that is too high.

Plotting the two against downstream accuracy across a sweep of hyperparameters produces a curve rather than a point cloud: for a fixed alignment, better uniformity tends to accompany better accuracy, and the frontier is the set of achievable trade-offs. The curve’s shape is the useful object. A method that moves the frontier outward is better than one that moves along it.

The temperature reappears

The contrastive loss with temperature τ corresponds to a particular weighting of the two terms, and changing τ moves the operating point along the trade-off curve. A smaller τ emphasizes alignment, because the sharp softmax concentrates gradient on the positive pair and on the hardest negatives. A larger τ emphasizes uniformity, because the flat softmax spreads gradient over many pairs and pushes the whole space apart.

This is a restatement of the temperature’s effect from InfoNCE and the Temperature Parameter, and it explains why temperature and batch size are usually tuned together: both change the position on the same curve.

Where the decomposition stops being exact

The asymptotic equivalence assumes the negative distribution covers the sphere and that the number of negatives is large. With a batch of a few hundred, the correspondence between the contrastive loss and the weighted sum of alignment and uniformity is loose, and a training run can reduce the loss without reducing either term in the way the decomposition predicts. The terms are still useful as diagnostics but not as an exact accounting.

The decomposition also assumes a uniform negative distribution, which supervised contrastive losses violate by construction. When labels determine which pairs are positives, the uniformity term is computed over a set that excludes same-class pairs, so the space is only required to be uniform across classes, not within them. The resulting geometry is a set of tight clusters spread over the sphere, and the within-cluster spread is not controlled by either term.

Relationship to other geometric accounts

Uniformity is one way to prevent collapse, but not the only one. Methods that use a stop-gradient and a momentum target prevent collapse structurally rather than geometrically, as described in Collapse and Redundancy Reduction in Self-Supervised Training. A space trained that way can have poor uniformity while still being useful, because the collapse was prevented by asymmetry rather than by spreading. The two mechanisms are not equivalent and can be combined.

Uniformity is also not the same as isotropy. A space can be uniform over a low-dimensional subspace, with all vectors lying near a great circle, and the pairwise distances will look spread within that subspace. Downstream tasks that need the remaining dimensions will fail. Measuring the effective rank of the representation matrix, for example by the entropy of its singular value spectrum, catches this case where the uniformity statistic does not.

Alignment and uniformity are also properties of the space at a given moment, and both change during training. Alignment typically improves early and then plateaus; uniformity often improves more slowly and can degrade if the learning rate is raised late. Monitoring both across training, rather than at the end, distinguishes a model that converged from one that is still moving in a direction the final measurement hides.

Why the two terms pull in opposite directions

The tension is structural rather than incidental. Alignment rewards agreement between the two views of the same input, and the cheapest way to achieve agreement is to reduce the distance between them. Since both representations are normalized, reducing the distance between a pair means moving them along the sphere toward each other, which increases the density of the region they occupy. Uniformity rewards separation, and the cheapest way to achieve separation is to move everything apart. Applied to the same parameters, the two gradients have opposite signs on the component that controls spread.

The balance point is not a fixed ratio. It depends on how much of the sphere the data can plausibly cover. A dataset with a small number of distinct concepts can be aligned tightly and still leave most of the sphere empty, so uniformity is easy to satisfy and the trade-off is mild. A dataset with many fine-grained concepts cannot afford tight alignment, because the concepts must be separated from each other as well as matched to their own views, and the achievable operating point moves toward uniformity.

This is one reason the effective capacity of the representation matters. A wider space can satisfy both terms at once because there is more room for many tight clusters, while a narrow space forces the clusters to overlap and the two terms compete for the same directions. The width of the projection head, which is usually narrower than the encoder output, is therefore a hyperparameter that acts on this trade-off directly rather than only on the loss arithmetic.

Practical consequences of the decomposition

Because the two terms can be computed on fixed data without training, they can be used to compare encoders trained by entirely different methods. Two models from different papers, evaluated on the same held-out set with the same positive-pair definition, can be placed on the same alignment-uniformity plane, and their positions say something about why one transfers better than the other. The comparison requires the positive-pair definition to be held fixed, since alignment is defined relative to a chosen notion of “same input” and that choice is the augmentation policy.

The decomposition also predicts where a downstream task will fail. A task whose classes are distinguished by a factor that the augmentations removed will sit in a region where alignment is high but the relevant variation has been collapsed, and no amount of uniformity helps. A task whose classes are rare will sit where uniformity is locally poor, because the rare classes occupy small regions that the average uniformity term barely sees. Both failures look like low accuracy but have different remedies, and separating the terms is what makes the distinction visible.

One caution about the Gaussian potential form. The exponent’s coefficient controls how quickly the potential decays with distance, and it interacts with the normalization convention. Implementations that use a different coefficient produce uniformity values on a different scale, so comparing the raw number across codebases is meaningless unless the convention matches. Comparing positions within one implementation is valid; comparing absolute values is not.

References

  1. Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere — Wang & Isola, ICML 2020
  2. Representation Learning with Contrastive Predictive Coding — van den Oord, Li & Vinyals, 2018
  3. Understanding Dimensional Collapse in Contrastive Self-supervised Learning — Jing, Vincent, LeCun & Tian, ICLR 2022