Representation Learning & Embedding Geometry

Let an encoder $f_\theta$ map each input $x$ to a vector $z$. If a simple predictor $g_\psi$ can recover a useful factor from $z$, the representation has made that factor easy to extract. For contextual word representations, the vector for token $k$ can depend on the whole sentence $t_1,\ldots,t_N$:

Here $\hat a$ is the predicted value of a factor or label $a$, $g_\psi$ is a simple extractor such as a linear classifier (a "probe"), and $h^{\mathrm{LM}}_{k,j}$ is the bidirectional language model's state for token $k$ at layer $j$, with $j=0$ the context-free token embedding. Peters et al. combine the $L+1$ layers with task-specific softmax weights $s_j^{\mathrm{task}}$ and a scale $\gamma^{\mathrm{task}}$. Because every $h^{\mathrm{LM}}_{k,j}$ with $j\ge1$ depends on the whole sentence, the same word receives different vectors in different sentences.

Similarity as a dot product (often cosine)

A common choice is cosine similarity, which is just a dot product of normalized vectors:

Normalization matters because it makes similarity depend on direction instead of length.

Contrastive learning (InfoNCE)

Suppose we have positive pairs $(x_i, y_i)$: two views of the same thing, such as two crops of one image or an image and its caption. Encode them with the same encoder ($\theta'=\theta$) or with two separate encoders:

InfoNCE treats $(z_i, u_i)$ as the correct match among a batch of candidates:

• Each term is a cross-entropy loss for picking the true partner $u_i$ out of the $N$ candidates in the batch; the other $N-1$ are the negatives.
• $\tau > 0$ is the temperature. Smaller $\tau$ sharpens the softmax, so most of the gradient comes from the negatives most similar to $z_i$.
• For independent pairs drawn from the joint distribution, with negatives drawn from the marginal distribution, van den Oord et al. (2018) bound the mutual information between the views by $I \ge \log N - \mathbb E[\mathcal L]$. The expectation is over sampled batches; the bound does not apply to every individual batch loss.

Checking for anisotropy

Two quick checks on a set of $N$ normalized embeddings $\hat z_i$ with mean $\bar z=\frac1N\sum_i\hat z_i$. First, the length $\lVert\bar z\rVert$: near $0$ when directions are spread out, near $1$ when all vectors point roughly the same way. Second, the covariance

If a few of its eigenvalues dominate, the vectors vary mainly along a few directions, which can dominate cosine comparisons. Subtracting the mean, or whitening with $C^{-1/2}$, changes this geometry.