Quadratic word embedding model (QWEM)

Karkada et al. (2025) Closed-Form Training Dynamics Reveal Learned Features and Linear Structure in Word2Vec-like Models

definitions & setup

Let $P(i,j):=P(j|i)P(i)$ denote the co-occurrence probability of words $i$ and $j$. Let $W\in {\mathbb{R}}^{V\times d}$ (where vocabulary size $V\gg d$) be the embedding matrix composed of row vectors $w_i$.

The standard Word2Vec (w2v) loss is parameterized as:

$$L_{\text{w2v}}(w,w^{\prime })={\mathbb{E}}_{i,j\sim P_{ij}}\left[{\Psi }_{ij}^{+}\log (1+e^{-w_i^T{w}_j^{\prime }})\right]+{\mathbb{E}}_{\begin{array}{c}i\sim P_i\\ j\sim P_j\end{array}}\left[{\Psi }_{ij}^{-}\log (1+e^{w_i^T{w}_j^{\prime }})\right]$$

• ${\Psi }_{ij}^{+}$ and ${\Psi }_{ij}^{-}$ are w2v preprocessing hyperparameters (these two degrees of freedom allow the loss to encode many engineering tricks used in preprocessing)
• the first term encourages large $w_i^T{w}_j^{\prime }$ for co-occurring pairs $(i,j)$
• the second term encourages small $w_i^T{w}_j^{\prime }$ for random pairs $(i,j)$
• here, $\log (1+e^{-x})=-\log (\text{sigmoid}(x))$

quadratic approximation (QWEM)

Using the second-order Maclaurin series approximation $\ln (1+e^x)\approx \ln 2+\frac{1}{2}\left(x+\frac{x^2}{4}\right)$, we define the QWEM loss under symmetric factorization ($w=w^{\prime }$):

$$L_w:={\mathbb{E}}_{i,j\sim P_{ij}}\left[{\Psi }_{ij}^{+}\left(\frac{(w_i^T{w}_j{)}^2}{4}-w_i^T{w}_j\right)\right]+{\mathbb{E}}_{\begin{array}{c}i\sim P_i\\ j\sim P_j\end{array}}\left[{\Psi }_{ij}^{-}\left(\frac{(w_i^T{w}_j{)}^2}{4}+w_i^T{w}_j\right)\right]$$

Define the optimization target for QWEM as $M_{ij}^{\ast }$:

$$M_{ij}^{\ast }:=\frac{{\Psi }_{ij}^{+}{P}_{ij}-{\Psi }_{ij}^{-}{P}_i{P}_j}{\frac{1}{2}({\Psi }_{ij}^{+}{P}_{ij}+{\Psi }_{ij}^{-}{P}_i{P}_j)}$$

Let the equivalence class under orthogonal transformation be:

$$\text{REquiv}(W):=\{WU\mid U\text{ is orthogonal}\}$$

Thm. QWEM as matrix factorization

Define the weighting term $G_{ij}:={\Psi }_{ij}^{+}{P}_{ij}+{\Psi }_{ij}^{-}{P}_i{P}_j$. Assume:

1. Symmetry: ${\Psi }_{ij}^{+}={\Psi }_{ji}^{+}$ and ${\Psi }_{ij}^{-}={\Psi }_{ji}^{-}$, which implies $M^{\ast T}=M^{\ast }$. Consequently, $M^{\ast }$ is diagonalizable: $M^{\ast }=V^{\ast }\Lambda V^{\ast T}$.
2. Constant weighting: $G_{ij}=g\in \mathbb{R}$. Statement: Under these assumptions, QWEM is equivalent to unweighted matrix factorization:

$$L_w\equiv \frac{g}{4}\Vert W{W}^T-M^{\ast }{\Vert }_F^2+\text{const}$$

Furthermore, if the truncated eigenvalue matrix ${\Lambda }_{[:d,:d]}$ is positive semi-definite (psd), then by the Eckart-Young-Minsky Theorem: e

$$\arg \underset{w}{\min }L(w)=\text{REquiv}\left(V_{[:,:d]}^{\ast }{\Lambda }_{[:d,:d]}^{1/2}\right)$$

Proof:

$$\begin{array}{rl}{L}_w& =\sum _{i,j}\left\{P_{ij}{\Psi }_{ij}^{+}\left(\frac{(w_i^T{w}_j{)}^2}{4}-w_i^T{w}_j\right)+P_i{P}_j{\Psi }_{ij}^{-}\left(\frac{(w_i^T{w}_j{)}^2}{4}+w_i^T{w}_j\right)\right\}\\ & \\ & =\sum _{i,j}\left\{\frac{(w_i^T{w}_j{)}^2}{4}g-w_i^T{w}_j({\Psi }_{ij}^{+}{P}_{ij}-{\Psi }_{ij}^{-}{P}_i{P}_j)\right\}\\ & \\ & =\frac{g}{4}\sum _{i,j}{\left(w_i^T{w}_j-\frac{{\Psi }_{ij}^{+}{P}_{ij}-{\Psi }_{ij}^{-}{P}_i{P}_j}{\frac{1}{2}g}\right)}^2+\text{const}\\ & \\ & =\frac{g}{4}\sum _{i,j}{\left((W{W}^T{)}_{ij}-M_{ij}^{\ast }\right)}^2+\text{const}\\ & \\ & =\frac{g}{4}\Vert W{W}^T-M^{\ast }{\Vert }_F^2+\text{const}■\end{array}$$

extensions & dynamics

• Relaxing Assumptions.
  • If we relax the constant $G_{ij}$ assumption, the loss becomes a weighted matrix factorization: $L(w)=\frac{1}{4}{\sum }_{i,j}{G}_{ij}{\left((W{W}^T{)}_{ij}-M_{ij}^{\ast }\right)}^2+\text{const}$
  • We can still find $\arg {\min }_{w}L(w)$ if ${\Psi }^{+},{\Psi }^{-}$ are symmetric and $G$ is a rank-1 matrix
• Optimization Dynamics.
  • Minimizing $L_{\text{QWEM}}(w)$ under Gradient Descent follows Saxe et al. dynamics. For an aligned initialization, the optimization problem reduces to growing the effective singular values

conclusions

• Linguistic Punchline: Natural language contains linear semantic structure within its co-occurrence statistics.
• ML Punchline: $\exists$ an explicit mathematical equivalence between self-supervised contrastive learning and supervised matrix factorization.

disclaimer: this note was mostly transcribed by Gemini