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.

Dimensional collapse and effective rank

Whitening assumes the variance is there to rebalance. Sometimes it is not. Dimensional collapse is when embeddings, without degenerating to a single point, occupy a subspace far smaller than d: a few large eigenvalues and a long tail of near-zeros. You paid for 768 dimensions and got the expressiveness of 30. The diagnostic is the singular-value spectrum, summarised as an entropy-based effective rank — RankMe:

p_k = σ_k / Σ_j σ_j        RankMe = exp( -Σ_k p_k log p_k )

A healthy d = 768 encoder scores in the hundreds; a collapsed one in the tens. It needs no labels and no downstream task, and it catches a failure the training curve hides: a model can drive InfoNCE down nicely while confining itself to a thin slab. Plot the log-spectrum — a cliff means collapse, a gentle decay means health.

If negatives exist mainly to prevent collapse, attack collapse directly and skip them, escaping the log N treadmill. Barlow Twins pushes the cross-correlation matrix C: [d, d] between two views’ embeddings toward the identity: diagonal for invariance, off-diagonal for redundancy reduction. VICReg splits that into an invariance MSE, a hinged variance term max(0, γ - sqrt(Var(z_j) + ε)) keeping every dimension alive, and a covariance term. Uniformity was never the point — rank preservation was.

Linear probing: measuring decodability, not information

The standard capability measure is a linear probe: freeze the encoder, fit only W: [d, C] on labelled data, and report accuracy. The restriction to a linear head is the whole point. Because I(z; y) ≤ I(x; y) by the data-processing inequality, a frozen encoder can never create information — so what probe accuracy varies with is not how much information is present but how linearly decodable it is.

Probing is thus the empirical counterpart of everything above: alignment puts same-class points near each other, healthy effective rank means the directions a classifier needs exist, whitening removes the offset that makes margins awkward. Probe accuracy rising while pre-training loss is flat means the geometry improved, not the information content. Two cautions: decodability says the information is readable, not that any downstream computation reads it; and an over-capacity probe learns the task instead of reading the representation.

What this means for small models on CPU

Every constraint here bites harder on a CPU-scale stack. The log N ceiling is the big one: if memory limits you to batches of 32, contrastive fine-tuning certifies at most log 32 = 3.47 nats, so borrow from the large-batch toolkit — a memory queue of stale embeddings, hard negatives mined offline, or a non-contrastive objective independent of N. And compare the implied bound log N - L across runs, never the raw loss.

On the inference side the wins are cheap and post-hoc. Whitening plus truncation is one small matrix multiply and often recovers more retrieval quality than swapping encoders. Effective rank is a two-minute SVD on a few thousand cached vectors, and it tells you whether a smaller d is safe: if a 768-dimensional model has effective rank 60, truncating to 256 costs almost nothing while cutting index memory by two-thirds — the logic Matryoshka embeddings train for deliberately.