Carl Allen: Homepage

Disentangling Disentanglement: how VAEs learn Independent Components

2024-12-15T20:39:00+00:00

This post summarises Unpicking Data at the Seams: VAEs, Disentanglement and Independent Components (Allen, 2024), which explains why disentanglement arises in generative latent variable models.

Disentanglement is an intriguing phenomenon observed in generative latent variable models, such as Variational Autoencoders (VAEs) (our focus), Generative Adversarial Networks (GANs) and latent Diffusion models. Although disentanglement has not been rigorously defined, it refers to when semantically meaningful factors of the data map to distinct dimensions in latent space. This allows, for example, images to be generated that vary by a particular feature by changing a single latent dimension (Figure 1 (left), or this video on disentanglement in GANs). It also allows apparent “embedding arithmetic”, e.g. the vector difference between embeddings of two images that vary in style can be added to that of a third image to transform its style (Figure 1, right).

While disentanglement is often associated with particular model families whose popularity may ebb and flow (e.g. VAEs, $\beta$-VAEs, GANs), we show that the phenomenon itself relates to the latent structure of the data and is more fundamental than any model that may expose it.

Figure 1. Illustrating disentanglement: (left) Each row of images varies by a single semantic feature, e.g. hair style or orientation. Disentanglement means that (i) such images can be generated by varying a _single_ latent dimension, i.e. by traversing the latent space in "axis-aligned" directions (parallel to standard basis vectors); and (ii) each latent variable affects a distinct semantic feature rather than several or all. (right) Changing the style of an image in diffusion models by disentangled features. (Images from Voynov & Babenko (2020) and Wu et al. (2023))

$\bullet$ Motivation: Disentanglement is particularly intriguing because models such as VAEs are not intentionally designed to achieve it, and it even occurs in settings impossible for other methods (i.e. isotropic Gaussian prior).¹$^,$² Thus understanding disentanglement may provide new insights into how models learn and what is learnable. More broadly, the ability to separate out independent aspects of the data is relevant to many areas of machine learning, and teasing its generative factors apart may offer fundamental insights into the data itself.

Machine learning has made incredible breakthroughs in recent years but our theoretical understanding lags notably behind, which may have serious implications as models increasingly reach the public sphere. Thus, a fundamental motivation for this work is to unravel some of the mystery of machine learning. We re-emphasise that disentanglement is not a model feature, rather a property of the data that a model may or may not expose and so is of more general interest. That said, autoencoding the data is often the first step in latent diffusion models, so understanding the latent space is relevant to current state of the art models.

$\bullet$ Approach: Recent works suggest that disentanglement arises in VAEs because commonly used diagonal posterior covariance matrices promote column-orthogonality in the decoder’s Jacobian.³$^,$⁴ In empirical support, columns of the decoder Jacobian are indeed found to be closer to orthogonal in VAEs with diagonal covariance than those with full covariance or general auto-encoders;³ and directly encouraging column-orthogonality with a constraint is shown to induce disentanglement.⁴

Building on this, Allen (2024): (A) further clarifies the connection between diagonal covariance and Jacobian orthogonality and (B) explains how disentanglement follows, showing that it amounts to factorising the data distribution into statistically independent components.

$$ \begin{equation} \text{diag. posterior covariance} \quad \overset{A}{\Rightarrow}\quad \text{column-orthog. Jacobian} \quad \overset{B}{\Rightarrow}\quad \text{disentanglement} \end{equation} $$

$\bullet$ Problem set-up: As in a VAE (or GAN), we assume that data $x\in\mathcal{X} \subseteq\mathbb{R}^n$ are generated by a latent variable model: independent latent variables are sampled $z\in\mathcal{Z} =\mathbb{R}^m$, $p(z) = \prod_ip(z_i)$, and mapped stochastically to the data space by sampling $x\sim p(x\mid z)$. The latent prior is commonly assumed to be isotropically Gaussian $p_\theta(z_i) = \mathcal{N}(z_i,0,1)$. We also assume $p_\theta(x\mid z)$ is Gaussian. [Notation: unsubscripted probability distributions denote the ground truth; subscripts indicate modelled distribution.]

TL;DR:

Disentanglement provably occurs in the linear case, $p_\theta(x\mid z) = \mathcal{N}(x; Dz,\sigma^2I),\ D\in\mathbb{R}^{n\times m}$, corresponding to probabilistic PCA (PPCA). Of the infinite set of known solutions, a VAE with diagonal posterior covariance finds only those in which latent dimensions $z_i$ map to independent factors of variation in the data distribution, which defines disentanglement.
Surprisingly, the rationale for the linear case (described later in this post) extends to non-linear VAEs with diagonal posterior covariance. The latter property encourages columns of the decoder’s Jacobian to be (approximately) orthogonal, which in turn means independent latent variables $z_i$ pass through the decoder and emerge (in $\mathcal{X}$) as statistically independent components that factorise the full push-forward distribution over the decoder-defined manifold.
- i.e. diagonal covariances cause the decoder to map independent factors in the latent space to independent components in the data space.
A VAE’s objective is maximised if the model distribution matches that of the data, hence if the data distribution factorises into independent components then the model distribution must similarly factorise. So independent factors of the data align with independent factors of the model, which (approximately) align with latent variables $z_i$ (from 2.) and so independent factors of $p(x)$ map to distinct latent variables of the model, i.e. $p(x)$ is disentangled.

$$ \begin{equation} \text{diag. posterior covariance} \quad \overset{A}{\Rightarrow}\quad \text{column-orthog. Jacobian} \color{lightgray}{\quad \overset{B}{\Rightarrow}\quad \text{disentanglement}} \end{equation} $$

(A) From Diagonal Covariance to Jacobian Orthogonality

The VAE fits a latent variable model $p_\theta(x) =\int_z p_\theta(x\mid z)p_\theta(z)$ to the data distribution $p(x)$ by maximising the Evidence Lower Bound (ELBO),

\[\ell(\theta, \phi) \quad =\quad \int p(x) \int q_\phi(z\mid x) \ \{\ \log p_\theta(x\mid z) \,-\, \beta \log \tfrac{q_\phi(z\mid x)}{p_\theta(z)} \ \}\ dz dx\ ,\]

where the ELBO has $\beta=1$ and $\beta>1$ is found to improve disentanglement. We assume common VAE assumptions:

$p_\theta(x\mid z) =\mathcal{N}(x;\,d(x),\sigma^2)\quad$ with decoder $d\in\mathcal{C}^2$ (injective) and fixed variance $\sigma^2$;
$q_\phi(z\mid x)=\mathcal{N}(z;\,e(x),\Sigma_x)\quad$ with encoder $e$ and learned variance $\Sigma_x$; and
$p_\theta(z)\quad\ \ \ =\mathcal{N}(z;\,0,I)\quad$ where $z_i$ are independent with $p_\theta(z_i)=\mathcal{N}(z_i;0,1)$.

Maximising the ELBO = maximum-likelihood$^{++}$: Maximising the likelihood $\int p(x)\log p_\theta(x)$ minimises the KL divergence between the data and model distributions, but this is often intractable for a latent variable model. Maximising the ELBO minimises the KL divergences between $p(x)q_\phi(z\mid x)$ and $p_\theta(x)p_\theta(z\mid x)\doteq p_\theta(x\mid z)p_\theta(z)$, aligning two models of the joint distribution.

The ELBO is optimised when $\Sigma_x$ relates to the Hessian of $\log p_\theta(x\mid z)$ (Opper & Archambeau) hence under certain conditions⁵ when:

\[\begin{equation} \Sigma_x^{-1} \ \ \overset{O\&A}{=}\ \ I - \mathbb{E}_{q(z\mid x)}[\tfrac{\partial^2\log p_\theta(x\mid z)}{\partial z_i\partial z_j}] \ \ \approx\ \ I + \mathbb{E}_{q(z\mid x)}[\tfrac{1}{\beta\sigma^2}J_z^\top J_z]\ , \tag{1}\label{eq:one} \end{equation}\]

where $J_z$ is the Jacobian of $d$ evaluated at $z$. For this to hold with diagonal $\Sigma_x$, columns of $J_z$ must be (approximately) orthogonal, $\forall z$.⁶ If $J_z=U_zS_zV_z^\top$ is the singular value decomposition (SVD), then $J_z$ is column-orthogonal $\quad$ iff $\ \ J_z^\top J_z = V_zS_z^2V_z^\top$ is diagonal $\quad$ iff $\ \ V_z=I$. Hence, by the definition of the Jacobian, diagonal covariances mean that perturbations in latent space in standard basis vectors $z_i$ (right singular vectors of $J_z$) correspond to perturbations in data space in directions $\mathbf{u}_i$ (column $i$ of $U_z$, left singular vectors of $J_z$), with no effect in any other $\mathbf{u}_{j\ne i}$. Note that this singles out axis-aligned directions in the latent space as “special”, breaking the rotational symmetry implied by the Gaussian prior.

Take-away: the ELBO is maximised as approximate posterior covariances $\Sigma_x$ tend to optimal posterior covariances defined by derivatives of $\log p_\theta(x\mid z)$. Diagonal covariances do not necessarily imply column-orthogonality in the decoder Jacobian, but if they exist, the VAE looks for solutions where the Hessian is diagonal hence where columns of the Jacobian are ($\approx$) orthogonal.

$$ \begin{equation} {\color{lightgray}{\text{diag. posterior covariance} \quad \overset{A}{\Rightarrow}}}\quad \text{column-orthog. Jacobian} \quad \overset{B}{\Rightarrow}\quad \text{disentanglement} \end{equation} $$

(B) From Orthogonality to Statistical Independence

Understanding that diagonal posterior covariance promotes column-orthogonality in the decoder Jacobian, the question then is how does that geometric property leads to the statistical property of disentanglement. For that we consider the push-forward distribution defined by the decoder, supported over a $m$-dimensional manifold $\mathcal{M}_d\subseteq\mathcal{X}$.

A push-forward distribution describes the probability distribution of the output of a deterministic function given an input distribution. A VAE decoder maps samples of the latent prior $p_\theta(z)$ to the data space, where it defines a manifold $\mathcal{M}_d$ and push-forward distribution over it.

Linear Case

For intuition, we first consider the linear case $x=d(z)=Dz$, $D\in\mathbb{R}^{n\times m}$, the model considered in Probabilistic PCA (PPCA), which has a tractable MLE solution and known optimal posterior

\[\begin{equation} p_\theta(z\mid x) = \mathcal{N}(z;\, \tfrac{1}{\sigma^2}M D^\top x,\, M) \quad\quad\quad M = (I + \tfrac{1}{\sigma^2}D^\top D)^{-1} \tag{2}\label{eq:two} \end{equation}\]

Since $D$ is the Jacobian of $d$, the latter expression is in fact a special case of \eqref{eq:one}. Hence if $D=USV^\top$ is the SVD, the ELBO is maximised when $\Sigma_x=M,\ \forall x\in\mathcal{X}$ and so, for diagonal $\Sigma_x$, when $V=I$.

Figure 2. Illustration of linear decoder $d:\mathcal{X}\to\mathcal{Z}$. The latent Gaussian prior is mapped to a Gaussian distribution in the data space over a linear manifold.

The aim is to understand how that constraint affects how independent dimensions $z_i\in\mathcal{Z}$ and the probability density over them pass through the decoder. For a given point $z^*\in \mathcal{Z}$, let:

$\mathcal{Z^{(i)}}\subset\mathcal{Z}$ be lines passing through $z^*$ parallel to each standard basis vector $e_i$ (blue dashed lines in Fig. 2), and $\mathcal{M}_d^{(i)}\subset\mathcal{M_d}$ be their images under $d$ that follow $\mathbf{u}_i$, the left singular vectors of $D$ (red dashed lines); and
$x^{(U)}=U^\top x=U^\top Dz^*$ be $x$ in the basis defined by columns of $U$⁷

Crucially, the Jacobian of the map from $z$ to $x^{(U)}$ is the diagonal matrix $S$, i.e. $\tfrac{\partial x^{(U)}_i}{\partial z_j} =\{s_i\doteq S_{i,i} \text{ if }i=j; 0 \text{ o/w}\}$. It then follows that independent $z_i$ map to independent components $x^{(U)}_i$ since:

each $x^{(U)}_i$ varies only with a distinct $z_i$ by considering $\tfrac{\partial x^{(U)}_i}{\partial z_j}$ and so are independent;
the push-forward of $d$ restricted to $\mathcal{Z^{(i)}}$ has density $p(x^{(U)}_i) = \mid\! s_i\!\mid ^{-1}p(z_i)$ over $\mathcal{M}_d^{(i)}$; and
the full push-forward distribution is $p(Dz) = \mid\!D\!\mid ^{-1}p(z) = \prod_i \mid\! s_i\!\mid^{-1}p(z_i) = \prod _ip(x^{(U)}_i)$.

This shows that the push-forward distribution defined by the decoder factorises as a product of independent univariate distributions $p(x^{(U)}_i)$, which each correspond to a distinct latent dimension $z_i$. Thus, if the data follows that generative process and so factorises (with factors determined by the SVD of ground truth $D$), then the ELBO is maximised when independent factors of the model align with those of the data and so p(x) is disentangled as a product of independent components that align with latent dimensions.

This may not seem a surprise in the linear case since it is known from the outset that the push-forward distribution is Gaussian and so factorises as a product of univariate Gaussians. However, neither that fact nor the linearity of $d$ was used at any stage.

Non-linear Case with Diagonal Jacobian

Figure 3. Illustration of non-linear decoder $d:\mathcal{X}\to\mathcal{Z}$. The latent Gaussian prior is mapped to a push-forward distribution in the data space over a non-linear manifold.

We now take an analogous approach for a general VAE ($x=d(z)$, $d\in\mathcal{C}^2$) with column-orthogonal decoder Jacobian. The Jacobian and its factors, $J_z=U_zS_zV_z^\top$, may now vary with $z$, however column-orthogonality restricts $V_z=I,\ \forall z\in \mathcal{Z}$ and $U_z$, $S_z$ are continuous w.r.t. $z$ (from $d\in\mathcal{C}^2$ and Papadopoulo & Lourakis (2006)).

As in the linear case, for any $z^*\in \mathcal{Z}$, we can define lines $\mathcal{Z^{(i)}}\subset\mathcal{Z}$ passing through $z^*$ parallel to the standard basis (blue dashed lines in Fig. 3), and their images under $d$, $\mathcal{M}_d^{(i)}\subset\mathcal{M_d}$, which are potentially curved sub-manifolds following (local) left singular vectors of $J_{z^*}$ (red dashed lines).

Considering $x=d(z)$ in the local $U$-basis (i.e. columns of $U_z$), denoted $x^{(U)}=U_z^\top x$, the Jacobian of the map from $z$ to $U_z^\top x^{(U)}$ is again diagonal $S_z$. Hence, as in the linear case, independent $z_i\in\mathcal{Z}$ map to independent components $x^{(U)}_i$ and:

$\{x^{(U)}_i\}_i$ are independent (by consideration of partial derivatives $\tfrac{\partial x^{(U)}_i}{\partial z_j}$);
the push-forward of $d$ restricted to $\mathcal{Z^{(i)}}$ has density $p(x^{(U)}_i) = \mid\! s_i\!\mid^{-1}p(z_i)$ over $\mathcal{M}_d^{(i)}$; and
the full push-forward satisfies $p(d(z)) = \mid J_z\mid ^{-1}p(z) = \prod_i \mid\! s_i\!\mid^{-1}p(z_i) = \prod _ip(x^{(U)}_i)$.

Thus, by the same argument as in the linear case, the distribution over the decoder manifold factorises as a product of independent univariate push-forward distributions $p(x^{(U)}_i)$, each corresponding to a distinct latent dimension $z_i$.

Putting everything together:

The ELBO is maximised if the model distribution fits the data distribution, so assuming that the data distribution has independent factors (by being generated under the considered model or otherwise) the model distribution must factorise similarly.
From part A, diagonal covariance matrices encourage the decoder’s Jacobian to be (approximately) column-orthogonal and, where so, the push-forward distribution over the model manifold factorises into components aligned with latent dimensions (from part B).
Thus the ELBO is maximised if independent components of the data distribution align with those of the model and if those of the model align with latent dimensions, thus the VAE aligns independent components that factorise the data distribution with latent dimensions, which defines disentanglement.
- (Identifiability depends on uniqueness of independent component distributions, analogous to uniqueness of singular values in the linear case).

Key insight: the above result hinges on the SVD of the Jacobian $J_z = U_zS_zV_z^\top$. By differentiability of $d$, the bases defined by columns of $U_z$ and $V_z$ (in $\mathcal{X}$ and $\mathcal{Z}$ resp.) are continuous in $z$ so basis vectors form continuous curves in each domain (these are linear in $\mathcal{Z}$ when $V_z=I$). By definition of the SVD, traversing a submanifold in one domain corresponds to traversing a corresponding submanifold in the other and the mapping between $x$ considered in the $U$-basis and $z$ considered in the $V$ basis has diagonal Jacobian given by $S_z$ - irrespective of how complex $d$ may be, whether a single linear layer or a thousand-layer CNN! As such, probability densities over submanifolds in $\mathcal{Z}$ map separably to their counterpart in $\mathcal{X}$, which is to say, independent components in $\mathcal{Z}$ map to independent components in $\mathcal{X}$.

Notes:

While we assume $d\in\mathcal{C^2}$, i.e. differentiability, the result also holds for continuous $d\in\mathcal{C^1}$ differentiable almost everywhere, e.g. ReLU networks.
We recommend reading the full paper for further details, such as:
- consideration of whether orthogonality is strictly necessary for disentanglement (the argument above shows it is sufficient);
- further consideration of model identifiability, i.e. up to what symmetries a VAE can identify ground truth generative factors; and
- explanation of the role of $\beta$ in a $\beta$-VAE, notably:
  - if $p_\theta(x\mid z)$ is of exponential family form then $\beta$ corresponds to Var$_\theta[x\mid z]$ (generalising $\sigma^2$ in the Gaussian case); and
  - $\beta$ acts a “glue” determining how close data points need to be (in Euclidean norm) for the model to treat them as “similar”, i.e. for their representations to merge.

Final remarks

We hope this post provides clearer insight into (i) what disentanglement means and how diagonal covariances, a seemingly innocuous design choice motivated by computational efficiency, leads to it; (ii) how statistically independent components of the data form “seams” running through the data manifold, which a VAE tries to “unpick”; and (iii) how understanding may transfer surprisingly well from linear to non-linear models.

This work throws up many interesting questions, e.g.:

if the VAE objective aims to disentangle, why is disentanglement not observed more reliably, e.g. as observed by Locatello et al. (2019)?
how does this relate to VAE extensions that enhance entanglement, such as FactorVAE and TC-VAE?
can this be used to automatically identify independent factors in the latent space of other models, e.g. GANs and latent diffusion models?
since this result is for continuous data domains, does it translate to discrete data such as text?
what does this insight into the structure of the data mean for other modelling paradigms, e.g. supervised, semi-supervised, self-supervised and reinforcement learning?

If any of these or related questions are of interest and you would like to collaborate, please free to get in touch (by email, twitter or bluesky).

Thanks for reading! We welcome any constructive feedback or discussion on this post on bluesky (I have pinned a link @carl-allen.bsky.social).

Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations; Locatello et al. (2019) ↩
Variational Autoencoders and Nonlinear ICA: A Unifying Framework; Khemakhem et al. (2020) ↩
Variational Autoencoders Pursue PCA Directions (by Accident); Rolinek et al. (2019) ↩ ↩²
On Implicit Regularization in β-VAEs; Kumar \& Poole (2020) ↩ ↩²
When second derivatives of the decoder are small almost everywhere, e.g. as in ReLU networks (see Abhishek & Kumar, 2020). ↩
As previously suggested more additional approximation steps and assumptions.³$^,$⁴ ↩
Note that we use slightly different notation to the original paper. ↩

‘Analogies Explained’ … Explained

2019-07-01T14:44:00+00:00

This post aims to present a `less math, more intuition’ overview of Analogies Explained: Towards Understanding Word Embeddings (ICML, 2019, Best Paper, hon mention), following the conference presentation. Target audience: ML, NLP, CL. [Skip to result]

Background

Word Embeddings

Word embeddings $\mathbf{w}=[w^{(1)}, w^{(2)}, … , w^{(d)}]$ are numerical vector representations of words (e.g. $\mathbf{w}_{\text{cat}}=[0.02, 0.29, …, -0.34])$, where each component $w^{(i)}$ can be thought of as capturing a semantic or syntactic feature of the word, and the full vector $\mathbf{w}$ acts as co-ordinates for the word in a high-dimensional “semantic space”.

Word embeddings can be created as a list of word co-occurrence statistics (or their low-rank factorisation); or learned by a neural network, such as Word2Vec (W2V), GloVe or a large language model. We shall see that these are closely related, and that the latter neural word embeddings learn statistics implied by the training algorithm, rather than them being explicitly defined.

Neural word embeddings (or simply embeddings) are useful across a wide variety of tasks involving natural language, from identifying word similarity or grammatical role, to assessing the sentiment in a customer review. Here, we focus specifically on word embeddings learned by Word2Vec and GloVe as the relative simplicity of their training algorithm and model architevture allows rigorous analysis.

Analogies

An intriguing property observed in trained word embeddings is that analogies are often solved simply by adding/subtracting them, e.g.:

“man is to king as woman is to queen”

can be solved using word embeddings by finding that closest to $\mathbf{w}_{king}-\mathbf{w}_{man}+\mathbf{w}_{woman}$ (omitting $man$, $woman$ and $king$ themselves), which often turns out to be $\mathbf{w}_{queen}$. This suggests that:

\[\begin{equation} \mathbf{w}_{king} - \mathbf{w}_{man} + \mathbf{w}_{woman} \approx \mathbf{w}_{queen}\, \tag{1}\label{eq:one} \end{equation}\]

or, in geometric terms, that word embeddings of analogies approximately form parallelograms:

While this seems intuitive, the fact that semantic relationships between words appear as geometric relationships between word embeddings is intriguing since word embeddings are not trained to achieve it! In practice, the parallelograms formed by word embeddings of analogies are not perfect:

This shows the (exact) parallelogram formed by $\mathbf{w}_{king}, \mathbf{w}_{man}, \mathbf{w}_{woman}$ and a fourth vertex ($\mathbf{w}_{king} {\small-} \mathbf{w}_{man} {\small+} \mathbf{w}_{woman}$) in the $xy$-plane with several word embeddings plotted relative to it. We see that (i) the embedding of $queen$ is not at the fourth vertex, but is the closest embedding to it; and (ii) that embeddings of related words, e.g. $prince$ and $lord$, are close relative to random words.

To explain the relationship between word embeddings of analogies \eqref{eq:one}, we describe the gap (marked ‘?’) between $\mathbf{w}_{queen}$ and $\mathbf{w}_{king} {\small-} \mathbf{w}_{man} {\small+} \mathbf{w}_{woman}$; and show that it is small, often smallest, for the word that we expect to complete the analogy.

We begin by considering what W2V embeddings learn.

Word2Vec

W2V (SkipGram with negative sampling) generates word embeddings by training the weights of a 2-layer “neural network” to predict context words $c_j$ (i.e. words that fall within a context window of fixed size $l$) around each word $w_i$ (the target word) over a text corpus.

To reduce the computational cost of computing $p(c_j\!\mid\! w_i)$ over all $c_j$ in a dictionary $\mathcal{E}$ using a softmax function, sigmoid functions were trained to classify truly occurring context words from words chosen at random (negative samples).

Levy & Goldberg (2014) showed that under this regime, matrices $\mathbf{W}, \mathbf{C}$ (whose columns are word embeddings $\mathbf{w}_i, \mathbf{c}_j$) approximately factorise a matrix of shifted Pointwise Mutual Information (PMI):

\[\begin{equation} \mathbf{W}^\top\mathbf{C} \approx \textbf{PMI} - \log k\ , \end{equation}\]

where $k$ is the chosen number of negative samples and $\begin{equation} \textbf{PMI}_{i,j} = \text{PMI}(w_i, c_j) = \log\tfrac{p(w_i,\, c_j)}{p(w_i)p(c_j)} \end{equation}\,.$

The shift $\log k$ is an arbitary artifact of the algorithm, which we address in the paper (Sec 5.5, 6.8) but drop here to simplify. The resulting relationship $\mathbf{W}^\top\!\mathbf{C} \!\approx\! \textbf{PMI}$ implies that an embedding $\mathbf{w}_i$ can be seen as a low-dimensional projection of $\text{PMI}_i$, the $i^{th}$ row of the PMI matrix, or PMI vector.

Proving the embedding relationship of analogies

This suggests observed additive relationships between word embeddings \eqref{eq:one} follow if the projection from PMI vectors to word embeddings induced by the loss function is sufficiently linear and if equivalent relationships exists between PMI vectors, i.e.

\[\begin{equation} \text{PMI}_{king} - \text{PMI}_{man} + \text{PMI}_{woman} \approx \text{PMI}_{queen}\ ; \tag{2}\label{eq:two} \end{equation}\]

A least squares loss gives precisely such a linear projection and we conjecture that W2V and Glove are sufficiently linear to preserve the relative spatial arrangement of vectors.

To now prove that relationship \eqref{eq:two} arises between PMI vectors of an analogy, we translate it into a relationship based on paraphrases that equate to an analogy.

Paraphrases

Inuitively, when we say a word $w_*$ paraphrases a set of words $\mathcal{W}$, we mean that they are semantically interchangeable in the text. For example, wherever $king$ appears, we might instead see both $man$ and $royal$. Mathematically, we define the best paraphrase $w^*$ as the word that maximises the likelihood of the context words found around $\mathcal{W}$, such that the context word distribution around $w_*$ (denoted $p(\mathcal{E}\!\mid\!w_*)$) is most similar to that around $\mathcal{W}$ ($p(\mathcal{E}\!\mid\!\mathcal{W})$), measured by Kullback-Leibler (KL) divergence. Although these distributions are discrete and unordered, we can picture this intuitively as:

Definition: $w_*$ paraphrases $\mathcal{W}$ if the paraphrase error ${\rho}^{\mathcal{W}, w_*}\!\in\! \mathbb{R}^{n}$ is (element-wise) small:
\[\begin{equation} {\rho}^{\mathcal{W}, w_*}_j = \log \tfrac{p(c_j|w_*)}{p(c_j|\mathcal{W})}\ , \quad c_j\!\in\!\mathcal{E}. \end{equation}\]

Paraphase error quantifies the similarity of the context distribution around any two words $w_1, w_2$ to that of any other word $w_*$ in terms of their PMI vectors:

The difference comprises the paraphrase error, which is small only when $w_*$ paraphrases $\{w_1, w_2\}$, and dependence error terms that do not depend on $w_*$. We formalise this as:

Lemma 1: For any word $w_*\!\in\!\mathcal{E}$ and word set $\mathcal{W}\!\subseteq\!\mathcal{E}$, $|\mathcal{W}|\!<\!l$, where $l$ is the context window size:
\[\begin{equation} \text{PMI}_* \ =\ \sum_{w_i\in\mathcal{W}} \text{PMI}_{i} \,+\, \rho^{\mathcal{W}, w_*} \,+\, \sigma^{\mathcal{W}} \,-\, \tau^\mathcal{W}\mathbf{1}\ . \end{equation}\]

This connects paraphrasing to PMI vector addition, as appears in \eqref{eq:two}, and is readily generalised to quantify the relationship between word sets. We say $\mathcal{W}$ and $\mathcal{W}_*$ paraphrase one another if the distributions of context words around them are similar (this reduces to the original definition if $\mathcal{W}_*$ contains a single word). Lemma 1 naturally extends to:

Lemma 2: For any word sets $\mathcal{W}, \mathcal{W}_*\!\subseteq\!\mathcal{E}$; $|\mathcal{W}|, |\mathcal{W}_*|\!<\!l$:
\[\begin{equation} \sum_{w_i\in\mathcal{W}_*} \text{PMI}_{i} \ =\ \sum_{w_i\in\mathcal{W}} \text{PMI}_{i} \,+\, \rho^{\mathcal{W}, \mathcal{W}_*} \,+\, (\sigma^{\mathcal{W}} - \sigma^{\mathcal{W}_*}) \,-\, (\tau^\mathcal{W} - \tau^{\mathcal{W}_*})\mathbf{1} \end{equation}\]

To see how this relates to paraphrases, let $\mathcal{W} \!=\! \{woman, king\}$ and $\mathcal{W}_* \!=\! \{man, queen\}$, whereby

\[\begin{equation} \text{PMI}_{king} - \text{PMI}_{man} + \text{PMI}_{woman} \ \approx\ \text{PMI}_{queen}\ , \end{equation}\]

if (a) $\mathcal{W}$ paraphrases $\mathcal{W}_*$ (i.e. $\rho^{\mathcal{W}, \mathcal{W}_*}$ is small) and (b) net statistical dependencies are small for words in $\mathcal{W}$ and $\mathcal{W}_*$ (i.e. $\sigma^{\mathcal{W}} - \sigma^{\mathcal{W}_*}$ and $\tau^\mathcal{W} - \tau^{\mathcal{W}_*}$).

Thus, subject to dependence errors, \eqref{eq:two} holds if $\{woman, king\}$ paraphrases $\{man, queen\}$.

This establishes that two paraphrases are a sufficient condition for the geometric relationship between word embeddings observed for an analogy. What remains then, is to show that an analogy “$w_a\text{ is to }w_{a^*}\text{ as }w_b\text{ is to }w_{b^*}$” implies that $\{w_b, w_{a^*}\!\}$ paraphrases $\{w_a, w_{b^*}\!\}$. In fact, we show an equivalence by reinterpreting paraphrases as word transformations.

Word Transformation

The paraphrasing of a word set $\mathcal{W}$ by a word $w_*$ can be interpreted as a semantic equivalence between $\mathcal{W}$ and $w_*$, denoted $\approx_p$. Alternatively, for a particular word $w$ in $\mathcal{W}$, a paraphrase indicates which words (the remainder of $\mathcal{W}$, denoted $\mathcal{W}^+$) that, when combined with $w$, make it “more like” – or transform it to – $w_*$. For example, the paraphrase of $\{man, royal\}$ by $king$ can be viewed as a word transformation from $man$ to $king$ by adding $royal$. The added words have the effect of narrowing the context or, more precisely, alter the distribution of context words around $w$ to more closely align with that of $w_*$. We can illustrate this as:

where the paraphrase acts as the “glue” in a relationship from $w$ to $w_*$. Thus, for $w\!\in\!\mathcal{W}$ and $\mathcal{W}^+ \!=\! \mathcal{W}\!\setminus\!\!\{w\}$, saying “$w_*$ paraphrases $\mathcal{W}$” is equivalent to saying “there exists a word transformation from $w$ to $w_*$ by adding $\mathcal{W}^+$”. Nothing changes other than the perspective.

Extending this to paraphrases between word sets $\mathcal{W}$, $\mathcal{W}_*\!$, we can choose $w \!\in\! \mathcal{W}$, $w_* \!\in\! \mathcal{W}_*$ and view the paraphrase as defining a relationship between $w$ and $w_*$ where $\mathcal{W}^+ \!=\! \mathcal{W}\!\setminus\!\!\{w\}$ is added to $w$ and $\mathcal{W}^- \!=\! \mathcal{W_*}\!\!\setminus\!\!\{w_*\}$ to $w_*$:

This is not (yet) a word transformation, as it lacks the notion of direction from $w$ to $w_*$. But, rather than viewing words in $\mathcal{W}^-$ as added to $w_*$, we consider them subtracted from $w$ (hence our naming convention):

Just as adding words narrows the context, subtracted words can be viewed as broadening the context.

Definition: there exists a word transformation from $w$ to $w_*$ with transformation parameters $\mathcal{W}^+$, $\mathcal{W}^-\subseteq\mathcal{E}$ iff $\ \{w\}\!\cup\!\mathcal{W}^+ \approx_\text{P} \{w_*\}\!\cup\!\mathcal{W}^-$.

Intuition

The intuition behind word transformations mirrors simple algebra: e.g. 5 is transformed to 8 by adding 3, but if 3 were somehow unavailable, then we could add 4 and subtract 1. With paraphrases alone, words can only be added to $w$ to transform it to $w_*$, and perhaps no word exists to describe such “semantic difference”. With word transformations, differences between words can also be used offering a far richer toolkit, e.g. the difference between $man$ and $king$ might be described crudely by, say, $royal$ or $crown$, but is perhaps better described by the difference between $woman$ and $queen$. More generally, transformation parameters ($\mathcal{W}^+, \mathcal{W}^-$) describe the difference between $w$ and $w_*$ or, one might say, how “$w$ is to $w_*$”.

Interpreting Analogies

We can now mathematically interpret the language of an analogy:

Definition (analogy): ”$w_a$ is to $w_{a^*}$ as $w_b$ is to $w_{b^*}\!$” iff there exist transformation parameters $\mathcal{W}^+\!, \mathcal{W}^-\!\subseteq\!\mathcal{E}$ that simultaneously transform both $w_a$ to $w_{a^*}$ and $w_b$ to $w_{b^*}$.

Each occurrence of “is to” in the analogy defintion is synonymous to a set of word transformation parameters and “as” implies their equality. Intuitively, this means that the semantic difference within each word pair, as now quantified, must be the same.

So, an analogy is a pair of word transformations with common parameters $\mathcal{W}^+, \mathcal{W}^-$, but what are those parameters? Fortunately, we need not search or guess. We show in the paper (Sec 6.4) that if an analogy holds, then any parameters that transform one of the word pairs, e.g. $w_a$ to $w_{a^*}$, must also transform the other pair. We can therefore choose $\mathcal{W}^+\!=\!\{w_{a^*}\!\}$, $\mathcal{W}^-\!=\!\{w_{a}\}$. These trivially transform $w_a$ to $w_{a^*}$ since $\{w_a, w_{a^*}\!\}$ exactly paraphrases $\{w_{a^*\!}, w_a\}$ (word order is irrelevant). For the analogy to hold, those parameters must then also transform $w_b$ to $w_{b^*}$, i.e. $\{w_b, w_{a^*}\}$ paraphrases $\{w_{b^*}, w_a\}$, that is:

\[\begin{equation} w_a \text{ is to } w_{a^*} \text{ as } w_b \text{ is to } w_{b^*}\!\!" \quad\iff\quad \{w_b, w_{a^*}\} \approx_\text{P} \{w_{b^*}, w_a\}. \tag{3}\label{eq:three} \end{equation}\]

This completes the chain:

analogies are equivalent to word transformations with common transformation parameters describing a common “semantic difference” \eqref{eq:three};
those word transformations define a readily identifiable paraphrase;
that paraphrase implies a geometric relationship between PMI vectors \eqref{eq:two} that depends on the paraphrase accuracy ($\rho$) and dependence error terms ($\sigma, \tau$); and
under a low-dimensional projection induced by the loss function, that geometric relationship is reflected between the word embeddings of an analogy \eqref{eq:one}.

Returning to an initial plot, we can now explain the “gap” in terms of paraphrase ($\rho$) and dependence $(\sigma, \tau$) error terms, and understand why it is variable but often smallest for words that complete the analogy. This concludes an end-to-end explanation for the geometric relationship between word embeddings observed for word analogies.

Extensions

“What the Vec? Towards Probabilistically Grounded Embeddings” extends the principles of this work to show: how W2V and Glove (approximately) capture other relationships such as relatedness and similarity; what certain embedding interactions correspond to; and, in doing so, how the semantic relationships of relatedness, similarity, paraphrasing and analogies mathematically inter-relate.
“Multi-relational Poincaré Graph Embeddings” [code] draws a comparison between analogies and relations in Knowledge Graphs to develop state-of-the-art representation models in both Euclidean and Hyperbolic space.

Several other works aim to theoretically explain the analogy phenomenon, in particular:

Arora et al. (2016) propose a latent variable model for text generation that is claimed inter alia to explain analogies, however strong a priori assumptions are made about the arrangement of word vectors that we do not require.
Gittens et al. (2017) introduce the idea of paraphrasing to explain analogies, from which we drew much inspiration, but they include several assumptions that fail in practice, in particular that word frequencies follow a uniform distribution rather than their true, highly non-uniform, Zipf distribution.
Ethayarajh et al. (2019) aim to show that word embeddings of analogies form parallelograms by considering the latter’s geometric properties. However, several strong assumptions are made that we do not require (e.g. all embedding must already be co-planar and embedding matrix $\mathbf{W}$ must be a scalar multiple of $\mathbf{C}$) and analogies are ultimately related to a statistical relationship “csPMI” that lacks semantic meaning.