Representation Learning & Embedding Geometry

Let an encoder map inputs to vectors. If a simple downstream predictor can recover a useful factor from the vector, the representation has made that factor easier to extract. For contextual word representations, the token vector can depend on the whole sentence:

$$\begin{aligned} z &= f_\theta(x)\in\mathbb R^d,\qquad \hat a = g_\psi(z),\\ \operatorname{ELMo}^{\mathrm{task}}_k &= \gamma^{\mathrm{task}}\sum_{j=0}^{L}s^{\mathrm{task}}_j\, h^{\mathrm{LM}}_{k,j}(t_1,\ldots,t_N). \end{aligned}$$

Here $a$ is a useful factor or label we hope is easy to read from $z$, $g_\psi$ is a simple extractor such as a classifier or probe, and $h^{\mathrm{LM}}_{k,j}$ is the bidirectional language-model state for token position $k$ at layer $j$. Peters et al. use task-specific softmax weights $s_j^{\mathrm{task}}$ and scale $\gamma^{\mathrm{task}}$ to combine the biLM layers, so the same token position can receive a different vector in a different sentence.

Similarity as a dot product (often cosine)

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

$$\mathrm{sim}(z, z') = \frac{z^\top z'}{\lVert z\rVert_2\,\lVert z'\rVert_2} = \hat z^\top \hat z'.$$

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: two crops of an image, two versions of a sentence, etc.). Encode them:

$$z_i = f_\theta(x_i), \qquad u_i = g_\theta(y_i).$$

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

$$\mathcal L = -\frac{1}{N} \sum_{i=1}^N \log \frac{\exp\big(\mathrm{sim}(z_i, u_i)/\tau\big)}{\sum_{j=1}^N \exp\big(\mathrm{sim}(z_i, u_j)/\tau\big)}.$$

• $\tau > 0$ is the temperature: smaller $\tau$ makes the softmax sharper (harder negatives, bigger gradients).
• In contrastive-learning setups, this kind of objective scores the intended pair above other candidates.

A useful geometry sanity check: anisotropy

Embeddings often become anisotropic (many points bunch along a few dominant directions). A practical check is the covariance of normalized embeddings:

$$C = \frac{1}{N}\sum_{i=1}^N (\hat z_i - \bar z)(\hat z_i - \bar z)^\top.$$

If a few eigenvalues dominate, the embedding cloud is highly directional, which can affect cosine-based comparisons. Normalization or whitening-style preprocessing can change this geometry.