Score-based Conditional Generation with Fewer Labeled Data by Self-calibrating Classifier Guidance

Song et al. [29] propose to model the transition from a known prior distribution $p_{T}({\bm{x}})$, typically a multivariate gaussian noise, to an unknown target distribution $p_{0}({\bm{x}})=p({\bm{x}})$ using the markov chain described by the following stochastic differential equation (SDE):

where $\bar{{\bm{w}}}$ is a standard Wiener process when the timestep flows from $T$ back to $0$, $s({\bm{x}},t)=\nabla_{\bm{x}}\log p_{t}({\bm{x}})$ denotes a time-dependent score function, and $f({\bm{x}},t)$ and $g(t)$ are some prespecified functions that describe the overall movement of the distribution $p_{t}({\bm{x}})$. The score function is learned by optimizing the following time-generalized denoise score matching (DSM) [31] loss

where $t$ is selected uniformly between $0$ and $T$, ${\bm{x}}_{t}\sim p_{t}({\bm{x}})$, ${\bm{x}}_{0}\sim p_{0}({\bm{x}})$, $s_{t}({\bm{x}}_{t}|{\bm{x}}_{0})$ denotes the score function of the noise distribution $p_{t}({\bm{x}}_{t}|{\bm{x}}_{0})$, which can be calculated using the prespecified $f({\bm{x}},t)$ and $g(t)$, and $\lambda(t)$ is a weighting function that balances the loss of different timesteps. In this paper, we use the hyperparameters from the original VE-SDE framework [29]. A more detailed introduction to learning the score function and sampling through SDEs is described in Appendix B.

In conditional SGMs, we are given labeled data $\{({\bm{x}}_{m},y_{m})\}_{m=1}^{M}$ in addition to unlabeled data $\{{\bm{x}}_{n}\}_{n=M+1}^{M+N}$, where $y_{m}\in\{1,2,\ldots,K\}$ denotes the class label. The case of $N=0$ is called the fully-supervised setting; in this paper, we consider the semi-supervised setting with $N>0$, with a particular focus on the challenging scenario where $\frac{M}{N+M}$ is small. The goal of conditional SGMs is to learn the conditional score function $\nabla_{\bm{x}}\log p({\bm{x}}|y)$ and then generate samples from $p({\bm{x}}|y)$, typically using a diffusion process as discussed in Section 2.1 and Appendix B.2.

One approach for conditional SGMs is classifier-free SGM [2, 6], which parameterizes its model with a joint architecture such that the class labels $y$ can be included as inputs. Classifier-free guidance [6], also known as CFG, additionally uses a null token $y_{\textsc{nil}}$ to indicate unconditional score calculation, which is linearly combined with conditional score calculation for some specific $y$ to form the final estimate of $s({\bm{x}}|y)$. CFG is a state-of-the-art conditional SGM in the fully-supervised setting. Nevertheless, as we shall show in our experiments, its performance drops significantly in the semi-supervised setting, as the conditional parts of CFG may lack sufficient labeled data during training.

Another popular family of conditional SGM is CGSGM. Under this framework, we decompose the conditional score function using Bayes’ theorem [29, 2]:

The $\log p(y)$ term can be dropped because it is not a function of ${\bm{x}}$ and is thus of gradient $0$. The decomposition shows that conditional generation can be achieved by an unconditional SGM that learns the score function $\nabla_{\bm{x}}\log p({\bm{x}})$ plus an extra conditional gradient term $\nabla_{\bm{x}}\log p(y|{\bm{x}})$.

The vanilla classifier-guidance (CG) estimates $\nabla_{\bm{x}}\log p(y|{\bm{x}})$ with an auxiliary classifier trained from the cross-entropy loss on the labeled data and learns the unconditional score function by the denoising score matching loss $\mathcal{L}_{\mathrm{DSM}}$, which in principle can be applied on unlabeled data along with labeled data. Nevertheless, the classifier within the vanilla CG approach is known to be potentially overconfident [17, 22, 21, 4] in its predictions, which in turn results in inaccurate gradients. This can mislead the conditional generation process and decrease class-conditional generation quality.

Dhariwal and Nichol [2] propose to address the issue by post-processing the term $\nabla_{\bm{x}}\log p(y|{\bm{x}})$ with a scaling parameter $\lambda_{CG}\neq 1$.

where $s({\bm{x}})=\nabla_{\bm{x}}\log p({\bm{x}})$ is the score function and $p(y|{\bm{x}};\boldsymbol{\phi})$ is the posterior probability distribution outputted by a classifier parameterized by $\boldsymbol{\phi}$. Increasing $\lambda_{CG}$ sharpens the distribution $p(y|{\bm{x}};\boldsymbol{\phi})$, guiding the generation process to produce less diverse but higher fidelity samples. While the tuning heuristic is effective in improving the vanilla CG approach, it is not backed by sound theoretical explanations.

Other attempts to regularize the classifier during training for resolving the issue form a promising research direction. For instance, CGSGM with denoising likelihood score matching [CG-DLSM; 1] presents a regularization technique that employs the DLSM loss below formulated from the classifier gradient $\nabla_{\bm{x}}\log p(y,t|{\bm{x}};\boldsymbol{\phi})$ and unconditional score function $s_{t}({\bm{x}})$.

where the unconditional score function $s_{t}({\bm{x}})$ is estimated via an unconditional SGM $s({\bm{x}}_{t},t;{\bm{\theta}})$. The CG-DLSM authors [1] prove that Eq. 2.2 can calibrate the classifier to produce more accurate gradients $\nabla_{x}\log p(y|{\bm{x}})$.

Robust CGSGM [11], in contrast to CG-DLSM, does not regularize by modeling the unconditional distribution. Instead, robust CGSGM leverages existing techniques to improve the robustness of the classifier against adversarial perturbations. Robust CGSGM applies a gradient-based adversarial attack to the ${\bm{x}}_{t}$ generated during the diffusion process, and uses the resulting adversarial example to make the classifier more robust.

Our proposed methodology draws inspiration from JEM [4], which shows that reinterpreting classifiers as energy-based models (EBMs) and enforcing regularization with related objectives helps classifiers to capture more accurate probability distributions. EBMs are models that estimate energy functions $E({\bm{x}})$ of distributions [15], which satisfies $\log p({\bm{x}})=-E({\bm{x}})+\log\int_{\bm{x}}\exp(E({\bm{x}}))d{\bm{x}}$. Given the logits of a classifier to be $f({\bm{x}},y;\boldsymbol{\phi})$, the estimated joint distribution can be written as $p({\bm{x}},y;\boldsymbol{\phi})=\frac{\exp(f({\bm{x}},y;\boldsymbol{\phi}))}{Z(\boldsymbol{\phi})}$, where $\exp(\cdot)$ means exponential and $Z(\boldsymbol{\phi})=\int_{{\bm{x}},y}\exp(f({\bm{x}},y;\boldsymbol{\phi}))\,d{\bm{x}}\,dy$. After that, the energy function $E({\bm{x}};\boldsymbol{\phi})$ can be obtained by

Then, losses used to train EBMs can be seamlessly leveraged in JEM to regularize the classifier, such as the typical EBM loss $\mathcal{L}_{\mathrm{EBM}}=\mathbb{E}_{p({\bm{x}})}\left[-\log p({\bm{x}};\boldsymbol{\phi})\right]$. JEM uses MCMC sampling for computing the loss and is shown to result in a well-calibrated classifier in their empirical study. The original JEM work [4] also reveals that classifiers can be used as a reasonable generative model, but its generation performance is knowingly worse than state-of-the-art SGMs.

We extend JEM [4] to connect the time-dependent classifier to the underlying distribution. In particular, we reinterpret the classifier as a time-dependent EBM. The interpretation allows us to obtain a time-dependent version of $p_{t}({\bm{x}})$ within the classifier, which can be used to obtain a classifier-internal version of the score function. Then, instead of regularizing the classifier by the EBM loss $-\log p_{t}({\bm{x}})$ like JEM, we propose to regularize by score function $\nabla_{\bm{x}}\log p_{t}({\bm{x}})$ instead.

Under the EBM interpretation, the energy function is $E({\bm{x}},t;\boldsymbol{\phi})=-\log\Sigma_{y}\exp(f({\bm{x}},y,t;\boldsymbol{\phi}))$, where $f({\bm{x}},y,t;\boldsymbol{\phi})$ is the output logits of the time-dependent classifier. Then, the internal time-dependent unconditional score function is $s^{c}({\bm{x}},t;\boldsymbol{\phi})=\nabla_{\bm{x}}\log\Sigma_{y}\exp(f({\bm{x}},y,t;\boldsymbol{\phi}))$, where $s^{c}$ is used instead of $s$ to indicate that the unconditional score is computed within the classifier. Then, we adopt the standard DSM technique in Eq. 2.1 to “train” the internal score function, forcing it to follow its physical meaning during the diffusion process. The resulting self-calibration loss can then be defined as

where ${\bm{x}}_{t}\sim p_{t}$, ${\bm{x}}_{0}\sim p_{0}$, and $s_{t}({\bm{x}}_{t}|{\bm{x}}_{0})$ denotes the score function of the noise centered at ${\bm{x}}_{0}$.

Fig. 2 summarizes the calculation of the proposed SC loss. Note that in practice, $t$ is uniformly sampled over $[0,T]$. After the self-calibration loss is obtained, it is mixed with the cross-entropy loss $\mathcal{L}_{\mathrm{CE}}$ to train the classifier. The total loss can be written as:

where $\lambda_{SC}$ is a tunable hyper-parameter. The purpose of self-calibration is to make the classifier to more accurately estimate the score function of the underlying data distribution, implying that the underlying data distribution itself is also more accurately estimated. As a result, the gradients of the classifiers are more aligned with the ground truth. After self-calibration, the classifier is then used in CGSGM to guide an unconditional SGM for conditional generation. Note that since our approach regularizes the classifier during training while classifier gradient scaling (Eq. 4) is done during sampling, we can easily combine the two techniques to enhance performance.

This section provides a comparative analysis of the regularization methods employed in DLSM [1], robust classifier guidance [11], JEM [4], and our proposed self-calibration loss.

In this work, we also explore the benefit of self-calibration loss in the semi-supervised setting, where only a small proportion of data are labeled. In the original classifier guidance, the classifiers are trained only on labeled data. The lack of labels in the semi-supervised setting constitutes a greater challenge to learning an unbiased classifier. With self-calibration, we better utilize the large amount of unlabeled data by calculating the SC loss with all data.

To incorporate the loss and utilize unlabeled data during training, we change the way $\mathcal{L}_{\mathrm{CLS}}$ is calculated from Eq. 8. As illustrated in Fig. 1, the entire batch of data is used to calculate $\mathcal{L}_{\mathrm{SC}}$, whereas only the labeled data is used to calculate $\mathcal{L}_{\mathrm{CE}}$. During training, we observe that when the majority is unlabeled data, the cross-entropy loss does not converge to a low-and-steady stage if the algorithm randomly samples from all training data. As this may be due to the low percentage of labeled data in each batch, we change the way we sample batches by always ensuring that exactly half of the data is labeled. Appendix C summarizes the semi-supervised training process of the classifier.

Note that even though the classifier is learning a time-generalized classification task, we can still make it perform as an ordinary classifier that classifies the unperturbed data by setting the input timestep $t=0$. This greatly facilitates the incorporation of common semi-supervised classification methods such as pseudo-labeling [16], self-training, and noisy student [32]. Integrating semi-supervised classification methodologies is an interesting future research direction, and we reserve the detailed exploration of this topic for future studies.

Following DLSM [1], we use a 2D toy dataset containing two classes to demonstrate the effects of self-calibration loss and visualize the training results. The data distribution is shown in Fig. 3(a), where the two classes are shown in two different colors. After training the classifiers on the toy dataset with different methods, we plot the gradients $\nabla_{\bm{x}}\log p(y|{\bm{x}})$ at minimum timestep $t=0$ estimated by the classifiers and compare them with the ground truth. Additional quantitative measurements of the toy dataset are included in Appendix E.

Figure 3 shows the ground truth classifier gradient and the gradients estimated by classifiers trained using different methods. Unregularized classifiers produce gradients that contain rapid changes in magnitude across the 2D space, with frequent fluctuations and mismatches with the ground truth. Such fluctuations can impede the convergence of the reverse diffusion process to a stable data point, leading SGMs to generate noisier samples. Moreover, the divergence from the ground truth gradient can misguide the SGM, leading to generation of samples from incorrect classes. Unregularized classifiers also tend to generate large gradients near the distribution borders and tiny gradients elsewhere. This implies that when the sampling process is heading toward the incorrect class, such classifiers are not able to “guide” the sampling process back towards the desired class. In comparison, the introduction of various regularization techniques such as DLSM, JEM, and the proposed self-calibration results in estimated gradients that are more stable, continuous across the 2D space, and better aligned with the ground truth. This stability brings about a smoother generation process and the production of higher-quality samples.

Table 1 and Fig. 4 present the performance of all methods when applied to various percentages of labeled data. CG-SC-labeled implies that self-calibration is applied only on labeled data whereas CG-SC-all implies that self-calibration is applied on all data.