Representation Learning Theory

What a representation actually is

Formally, an encoder f maps an input x to z = f(x) with z: [d], and a downstream task applies a head g from some restricted class G to get a prediction ŷ = g(z). Two properties matter. Sufficiency: z retains the label information, I(z; y) ≈ I(x; y). Minimality: z discards everything else, so I(z; x) is as small as sufficiency permits.

That pair is the information-bottleneck framing, and it is necessary but not sufficient, because it says nothing about form. A perfectly sufficient 4096-bit hash of the input would be useless: nothing in a realistic G could read it. The real target is sufficiency plus a geometry your intended head class can decode cheaply.

Positives, negatives, and the hypersphere

Without labels you need another signal for what should be close. The standard device is a positive pair: two views x and x⁺ of the same content — two crops of an image, a query and the passage answering it. Everything else in the batch is a negative.

Encode both views and L2-normalise, so every vector lives on the unit hypersphere S^(d-1). There the dot product is cosine similarity, and squared distance and similarity are two views of one quantity:

||u - v||^2 = 2 - 2(u · v)   for ||u|| = ||v|| = 1

Normalising is not cosmetic: it removes the norm as a free parameter the loss could exploit, bounds the logits, and reduces the problem to one about angles — which is what makes the geometric analysis tractable.

InfoNCE and its log N ceiling

Given a batch of N pairs, InfoNCE treats ‘which of the N candidates is my partner?’ as an N-way classification problem, scored by a temperature-scaled softmax over similarities:

L_NCE = -(1/N) Σ_i log [ exp(z_i · z_i⁺ / τ)
                      -----------------------------
                      Σ_j exp(z_i · z_j⁺ / τ) ]

shapes:  Z: [N, d]  Z⁺: [N, d]  S = Z Z⁺^T: [N, N]
target:  the diagonal of S  →  cross-entropy with labels 0..N-1

That is one matrix multiply and one cross-entropy; the off-diagonal entries are in-batch negatives, free contrast since every other example had to be encoded anyway, at O(N^2 d) — negligible next to the encoder passes. What makes it principled is that its optimum is a variational lower bound on the mutual information between views, I(x; x⁺) ≥ log N - L_NCE: the Bayes-optimal scorer for the N-way task is the density ratio p(x⁺ | x) / p(x⁺), so driving the loss down certifies information. The sting is the ceiling — with a batch of N you can never certify more than log N nats, and log 256 = 5.55 nats is only about 8 bits. Hence the appetite for huge batches, memory banks and momentum queues; past saturation the loss stops measuring information and measures geometry.

A worked numeric example

Take N = 4 and an anchor whose cosine similarities against the four candidates are [0.60, 0.45, 0.40, 0.30], the first being the true positive. At τ = 0.1 the logits are [6.0, 4.5, 4.0, 3.0]:

exp:   403.43   90.02   54.60   20.09    sum = 568.14
p(pos) = 403.43 / 568.14 = 0.7101
loss   = -ln(0.7101)      = 0.3424 nats
bound  = log 4 - 0.3424   = 1.3863 - 0.3424 = 1.044 nats  (≈ 1.51 bits)

Now sharpen to τ = 0.05, doubling the logits to [12, 9, 8, 6]. Softmax mass on the positive jumps to 0.9341, the loss falls to 0.0682, and the certified bound rises to 1.318 of a possible 1.386 nats. Nothing about the encoder changed — only the temperature. Two lessons: loss values are never comparable across temperatures or batch sizes, since both rescale L_NCE directly; and at N = 4 the 2-bit ceiling is so low that near-zero loss certifies almost nothing.

Alignment and uniformity on the hypersphere

The cleanest account of what InfoNCE optimises splits it into two geometric quantities. Alignment asks that positive pairs land close together; uniformity asks that the marginal distribution of embeddings spread out over the sphere:

L_align  = E_(x,x⁺) ||f(x) - f(x⁺)||^2
L_unif   = log E_(x,y ~ p, iid) exp(-2 ||f(x) - f(y)||^2)

As N → ∞ the InfoNCE numerator drives alignment and the denominator drives uniformity. The tension is real: alignment alone is minimised by the constant map (everything to one point, loss zero, information zero), uniformity alone by ignoring content entirely. Negatives supply the repulsive force, so a contrastive objective without them collapses. Temperature sets the balance — small τ weights the hardest negatives most heavily, buying uniformity at the cost of tolerance for true near-duplicates.

Anisotropy, and whitening as the fix

Theory says uniformity; measured spaces say otherwise. Take mean cosine similarity over random unrelated pairs of hidden states from a trained transformer: instead of the ≈ 0 you would expect on a 768-dimensional sphere, 0.4 to 0.8 is routine. The representations occupy a narrow cone. Blame token frequency and the softmax bottleneck of language-model training, which leave a few rogue dimensions of enormous variance dominating every dot product. The damage is lost resolution: if every pair scores 0.55 to 0.75, the signal is a 0.2-wide band on a constant offset, and thresholds transfer badly. But this is a mean problem — the space can be off-centre yet full rank — so a correction on the first two moments largely removes it:

μ = (1/n) Σ_i x_i          Σ = (1/n) (X - μ)^T (X - μ)   : [d, d]
Σ = U Λ U^T   (eigendecomposition)
W = U Λ^(-1/2)      : [d, d]
z = (x - μ) W        ⇒  Cov(z) = I

Centring kills the offset; Λ^(-1/2) equalises every direction so no rogue dimension dominates. Mean random-pair cosine typically drops from ~0.6 to ~0.0 and retrieval gains a few points, for one [d, d] matmul per query. Two caveats: Σ must be fitted on data resembling your corpus (one fitted on Wikipedia will not fix support tickets), and whitening amplifies the smallest eigendirections, mostly noise — so truncate to the top k components.