Learning to Be Guided: Shaping Latent Geometry for Classifier-Free Guidance

What does CFG actually add?

CFG amplifies the displacement between the conditional and unconditional predictions. For the Bayes-optimal predictors \(v_c^\star=\mathbb E[v\mid z_t,t,C=c]\) and \(v_u^\star=\mathbb E[v\mid z_t,t]\), substituting \(v=(z-z_t)/(1-t)\) turns that displacement into a difference of posterior means of the clean latent:

The correction is how far knowing the class moves the estimate of the clean latent. To trace its dependencies, we need a model where it can be computed in closed form.

When is the correction large, and when does it vanish?

Take \(K\) classes with Gaussian priors \(p(z\mid c)=\mathcal N(z;\mu_c,\sigma^2 I)\) sharing one intra-class variance \(\sigma^2\). Writing \(q_j=\Pr(C=j\mid z_t,t)\) for the noisy class posterior, Eq. 1 becomes

class contrast \(\mu_c-\mu_j\)posterior weight \(q_j\)intra-class variance \(\sigma^2\)

The correction sums displacements from each competitor mean toward the target mean, weighted by that competitor's posterior \(\class{cw}{q_j}\) (i.e., how confusable it remains at the noisy state \(z_t\)). Once the target class dominates the posterior (\(q_c\to 1\)), all competitor weights vanish, and so does the correction.

Which geometry keeps classes distinguishable under noise?

Step 2 says the correction lives as long as the posterior is uncertain. How soon it becomes certain depends on how far apart classes remain once noise is added. Since \(p(z_t\mid c)=\mathcal N\big(z_t;\,t\mu_c,\,(t^2\sigma^2+(1-t)^2)I\big)\), the squared standardized distance between two classes is

To separate geometry from scale, fix reference means \(\mu_c^{\text{ref}}\in\mathbb R^d\) with \(\mathbb E_C[\mu_C^{\text{ref}}]=0\) and \(\mathbb E_C\|\mu_C^{\text{ref}}\|^2=d\), and vary only how a fixed total variance \(d\) is split:

Raising \(\alpha\) moves variance from intra-class \((1-\alpha)d\) to inter-class \(\alpha d\), pushing centers apart and shrinking intra-class spread. This increases \(\delta^2_{cj}(t)\) at every \(t\), making classes more distinguishable under noise, so the target posterior concentrates earlier.

How do we turn this into a training objective?

To keep classes distinguishable under noise, we would ideally minimize the class uncertainty left at noise level \(t\), measured by the cross-entropy of the true noisy class posterior:

Holding the posterior function \(q_c\) fixed, the per-sample descent direction on the clean latent \(z\) is proportional to the oracle correction \(g_c^\star\):

To turn Eq. 5 into a practical training objective, we also need to choose how \(z_t\) is sampled. A natural choice follows the training trajectory: drawing \(t\) from a time distribution \(\rho\) gives \(\mathcal L_\rho=\mathbb E_{t\sim\rho}[\mathcal L_{\text{post}}(t)]\), with \(\rho\) ranging from the whole trajectory (uniform \(\rho\)) to clean latents only (\(\rho\) concentrated at \(t=1\)). Within the family of class geometries defined in Eq. 4, every \(\mathcal L_\rho\) is non-increasing in \(\alpha\), so all of them favor the same geometry: more separated and compact classes.

The objective is aligned with CFG at two levels: (1) its per-sample gradient follows CFG's correction (Eq. 6), and (2) minimizing any \(\mathcal L_\rho\) promotes the separated, compact geometry that strengthens CFG (Step 3). In practice, one then needs to choose \(\rho\) and approximate \(q_c\).