Learning Embeddings with Diffusion Models
The goal of this note is to understand what goes wrong when the embeddings of a discrete space are learned jointly with a continuous diffusion model.
The short version is the following. If the embeddings are fixed and injective, a diffusion model in embedding space is just a continuous way of modeling a discrete distribution. But if the embeddings are optimized by the same denoising loss, the target itself starts moving. The optimizer can then improve the loss by making the denoising problem easier, even if it destroys the representation of the original discrete distribution.
This gives three different pathologies:
- with an MSE denoising loss, embeddings want to collapse;
- with a cross-entropy posterior loss, embeddings want to separate so much that classification becomes trivial;
- with the posterior KL alone, an expressive model has zero loss for every embedding geometry, so the embeddings receive no principled signal.
So learnable embeddings can be useful, but only as a constrained coordinate system or inductive bias. They are not identified by the diffusion surrogate itself.
1. The setting
Let X \in \mathcal X = \{1,\ldots,K\} be sampled from the discrete data distribution p_0. An embedding map e_\phi assigns a vector e_i = e_\phi(i) \in \mathbb R^d to each category i. This turns p_0 into the embedded measure
\mu_\phi = \sum_{i=1}^K p_0(i)\delta_{e_i}.
If the embeddings are fixed and injective on the support of p_0, this is harmless. A perfect model of \mu_\phi, followed by the decoder \operatorname{Dec}_\phi(e_i)=i, recovers p_0. Indeed, if Y\sim \mu_\phi, then \mathbb P(\operatorname{Dec}_\phi(Y)=i)=\mathbb P(Y=e_i)=p_0(i).
This is an interpretation of embeddings: they are coordinates for the discrete objects. As such the simplest embeddings is the one-hot encoding. They may make the modeling problem smoother or easier for a finite neural network, just as word embeddings in language models let nearby vectors share statistical structure.
The difficulty starts when \phi is optimized by the diffusion loss itself. Then the target is no longer just p_0 but the moving measure \mu_\phi. The loss can decrease because the model learned the distribution better, but it can also decrease because \phi changed the distribution into something easier to denoise.
For the rest of the note, take a Gaussian corruption process
Z_t = a_t e_\phi(X) + \sigma_t \varepsilon,\qquad \varepsilon\sim \mathcal N(0,I),
where t is sampled from the training schedule. The precise schedule is not important for the argument, what matters is that positive-noise times have \sigma_t>0.
2. The MSE denoising loss collapses embeddings
The standard denoising objective is
\mathcal L_{\mathrm{MSE}}(\theta,\phi) = \mathbb E\left[\|e_\phi(X)-D_\theta(Z_t,t)\|^2\right].
This objective is usually justified by a simple projection identity. The real regression target of the denoiser is the posterior mean m_\phi(Z_t,t)=\mathbb E[e_\phi(X)\mid Z_t,t]. The corresponding posterior-mean objective would be
\mathcal J_{\mathrm{MSE}}(\theta,\phi) = \mathbb E\left[\|m_\phi(Z_t,t)-D_\theta(Z_t,t)\|^2\right].
For fixed embeddings, \mathcal L_{\mathrm{MSE}} and \mathcal J_{\mathrm{MSE}} differ only by a term that is constant in \theta, so one can drop it. This is the usual denoising trick. But if \phi is learned, the dropped term is no longer constant. It depends on \phi, and optimizing it is exactly what creates the collapse.
For fixed \phi, the optimal denoiser is D^\star_\phi(z,t)=\mathbb E[e_\phi(X)\mid Z_t=z,t]. Therefore the optimized loss is the posterior variance:
\min_\theta \mathcal L_{\mathrm{MSE}}(\theta,\phi) = \mathbb E\operatorname{Tr}\operatorname{Var}(e_\phi(X)\mid Z_t,t).
This identity explains the collapse. If all embeddings on the support of p_0 satisfy e_i=c, then e_\phi(X)=c almost surely. The conditional variance is zero and the constant denoiser D_\theta(z,t)=c achieves zero loss.
Since the loss is nonnegative, this is a global minimum. Under positive Gaussian noise and at least two classes with positive probability, it is essentially the only way to get exact zero loss: because the Gaussian components have full support, Z_t cannot reveal the class perfectly at finite embedding distances. Thus \operatorname{Var}(e_\phi(X)\mid Z_t,t)=0 forces e_\phi(X) to be almost surely constant. If some classes have p_0(i)=0, only the embeddings on the support matter.
The key point is:
\boxed{\text{MSE makes the denoising problem easier by allowing } e_\phi(X) \text{ itself to become nearly constant.}}
Once e_\phi(X) is constant, there is no denoising problem left. But the discrete distribution has not been learned, it has been erased by the representation.
Let m_\phi(z,t)=\mathbb E[e_\phi(X)\mid Z_t=z,t]. Then
\begin{aligned} \mathbb E\|e_\phi(X)-D_\theta(Z_t,t)\|^2 &= \mathbb E\|e_\phi(X)-m_\phi(Z_t,t)\|^2 \\ &\quad+ \mathbb E\|m_\phi(Z_t,t)-D_\theta(Z_t,t)\|^2. \end{aligned}
The cross term vanishes because \mathbb E[e_\phi(X)-m_\phi(Z_t,t)\mid Z_t,t]=0.
For fixed \phi, the first term does not depend on \theta, so minimizing the denoising loss over \theta is equivalent to minimizing the posterior-mean objective
\mathbb E\|m_\phi(Z_t,t)-D_\theta(Z_t,t)\|^2.
This is the usual score-matching/denoising trick: replace the intractable posterior mean target by the noisy clean sample target, because the difference is a constant.
But when \phi is also optimized, the first term is no longer constant. It is exactly
\mathbb E\|e_\phi(X)-m_\phi(Z_t,t)\|^2 = \mathbb E\operatorname{Tr}\operatorname{Var}(e_\phi(X)\mid Z_t,t).
The joint objective therefore rewards embeddings that reduce the intrinsic posterior variance. Collapsed embeddings make this variance zero. So the problem is not that the projection identity is false; it is that the term usually treated as a constant is constant only with respect to \theta, not with respect to \phi.
3. Cross-entropy pushes in the opposite direction
A common alternative for categorical data is to predict the clean class from the noisy embedding. Let p_\theta(i\mid z,t) be the model’s predicted categorical posterior. The cross-entropy loss is
\mathcal L_{\mathrm{CE}}(\theta,\phi) = \mathbb E[-\log p_\theta(X\mid Z_t,t)].
For fixed \phi, write p_\phi(i\mid z,t)=\mathbb P(X=i\mid Z_t=z,t). Then
\mathcal L_{\mathrm{CE}}(\theta,\phi) = \mathbb E_{Z_t,t}\left[ H(p_\phi(\cdot\mid Z_t,t)) + \mathrm{KL}\!\left( p_\phi(\cdot\mid Z_t,t)\,\middle\|\,p_\theta(\cdot\mid Z_t,t) \right) \right].
So the optimal classifier is p_\theta=p_\phi, and the optimized loss is
\min_\theta \mathcal L_{\mathrm{CE}}(\theta,\phi) = H_\phi(X\mid Z_t,t) = H(X)-I_\phi(X;Z_t\mid t).
Now optimizing \phi tries to minimize H_\phi(X\mid Z_t,t), or equivalently maximize the mutual information between the noisy embedding and the original class.
This is the reverse of the MSE pathology. If the embeddings collapse, Z_t contains no class information and the optimized cross-entropy is H(X). If the embeddings are sent far apart relative to the noise level, then Z_t almost determines X, the posterior p_\phi(\cdot\mid Z_t,t) becomes almost a Dirac mass, and the optimized cross-entropy tends to zero.
Thus:
\boxed{\text{CE does not collapse the embeddings, it makes classification trivial by separating them.}}
Without a scale constraint, this usually has no finite minimizer: the infimum is approached by sending pairwise distances to infinity. With a norm or covariance constraint, the problem becomes a packing or max-information problem. The minimizer may not be unique, and there is no reason for the resulting geometry to be semantic unless extra structure forces it to be.
Condition on Z_t=z and t. The conditional contribution to the cross-entropy is
\mathbb E[-\log p_\theta(X\mid z,t)\mid z,t] = -\sum_i p_\phi(i\mid z,t)\log p_\theta(i\mid z,t).
Add and subtract \sum_i p_\phi(i\mid z,t)\log p_\phi(i\mid z,t):
\begin{aligned} -\sum_i p_\phi(i\mid z,t)\log p_\theta(i\mid z,t) &= H(p_\phi(\cdot\mid z,t)) \\ &\quad+ \mathrm{KL}\!\left( p_\phi(\cdot\mid z,t)\,\middle\|\,p_\theta(\cdot\mid z,t) \right). \end{aligned}
After averaging over (Z_t,t), we get the decomposition above. For fixed \phi, the entropy term is constant in \theta. For learned \phi, it is not constant, and cross-entropy rewards embeddings that make the posterior entropy small.
4. The proper objectives alone do not learn embeddings
The previous sections suggest a natural correction: keep only the part of the surrogate that really fits the desired conditional object, and remove the term that was only constant when \phi was fixed.
For MSE, this means training the posterior mean directly:
\mathcal J_{\mathrm{MSE}}(\theta,\phi) = \mathbb E\left[\|m_\phi(Z_t,t)-D_\theta(Z_t,t)\|^2\right], \qquad m_\phi(Z_t,t)=\mathbb E[e_\phi(X)\mid Z_t,t].
This is the objective that the denoiser actually cares about. The standard denoising loss reduces to it only after dropping \mathbb E\operatorname{Tr}\operatorname{Var}(e_\phi(X)\mid Z_t,t), and that drop is legitimate only if \phi is fixed.
But \mathcal J_{\mathrm{MSE}} also shows the same non-identifiability in another form. With an expressive enough denoiser, for every \phi one can set D_\theta(z,t)=m_\phi(z,t), and then \min_\theta \mathcal J_{\mathrm{MSE}}(\theta,\phi)=0. So the proper posterior-mean objective removes the collapse incentive, but it does not by itself select a meaningful embedding geometry.
For cross-entropy, the analogous correction is the same one discussed above. If the problematic part of cross-entropy is the entropy term H(p_\phi(\cdot\mid Z_t,t)), why not train only the posterior KL?
That objective would be
\mathcal L_{\mathrm{KL}}(\theta,\phi) = \mathbb E_{Z_t,t}\, \mathrm{KL}\!\left( p_\phi(\cdot\mid Z_t,t)\,\middle\|\,p_\theta(\cdot\mid Z_t,t) \right).
This is the proper posterior-fitting objective. The standard cross-entropy reduces to it only after dropping H(p_\phi(\cdot\mid Z_t,t)), and that entropy is constant only when \phi is fixed. With an expressive enough classifier, for every admissible \phi we can set p_\theta(\cdot\mid z,t)=p_\phi(\cdot\mid z,t), and then \min_\theta \mathcal L_{\mathrm{KL}}(\theta,\phi)=0.
So the proper objectives remove the bad incentives, but they also remove the embedding signal at the exact expressive optimum. In a finite-capacity model, \phi may still move toward geometries that make m_\phi or p_\phi(\cdot\mid z,t) easier to approximate. That can be useful, but it is an inductive-bias effect, not an exact statistical principle that identifies the correct embedding geometry.
A way to phrase the issue is:
\boxed{\text{The diffusion/posterior surrogate can train a model for a given geometry, but it cannot by itself decide what the geometry should mean.}}
5. How to fix the joint-learning problem
There is no free fix that lets the diffusion denoising loss learn unconstrained embeddings and still preserves the right minimizers. One has to add structure.
A conservative recipe is to treat embeddings as a coordinate system:
- use fixed or pretrained embeddings;
- or constrain the learned codebook, for example by fixing norms, centering, whitening, or otherwise preventing collapse and scale blow-up;
- and keep an explicit decoder or reconstruction condition so token identity is preserved.
A cleaner joint objective separates the roles:
\mathcal L(\theta,\phi,\psi) = \mathcal L_{\mathrm{diff}}(\theta;\operatorname{sg}(\phi)) + \lambda\,\mathcal L_{\mathrm{repr}}(\phi,\psi),
where \operatorname{sg}(\phi) means that the diffusion loss does not update the embeddings. The representation loss can include a token-preserving decoder term such as
\mathcal L_{\mathrm{repr}}(\phi,\psi) = \mathbb E[-\log p_\psi(X\mid e_\phi(X))] + \text{geometry regularization}.
The first term prevents the embedding from forgetting the class. The regularization prevents degenerate geometry. Then the diffusion model is trained for the current geometry, while the geometry is learned by an objective that actually cares about token identity.
More ambitious approaches build reconstruction or consistency terms directly into the diffusion objective [2].
6. Bottom line
For fixed injective embeddings, diffusion in embedding space is a valid way to model p_0. For learned unconstrained embeddings, the standard losses become ill-posed:
\text{MSE} \Rightarrow \text{collapse},\qquad \text{CE} \Rightarrow \text{infinite separation}.
Therefore learnable embeddings can help only when they are controlled by constraints, reconstruction, likelihood, or consistency conditions that preserve the discrete identity of the data. Their value is mainly finite-capacity and optimization value: they can choose a geometry in which the posterior or score is easier for a neural network to approximate. They do not, by themselves, change the exact target distribution.
References:
[1] Sander Dieleman et al., “Continuous Diffusion for Categorical Data”, arXiv:2211.15089.
[2] Bac Nguyen et al., “Improving Vector-Quantized Image Modeling with Latent Consistency-Matching Diffusion”, arXiv:2410.14758.