BGDB: Bernoulli-Gaussian Decision Block with Improved Denoising Diffusion Probabilistic Models

TL;DR

This work tackles the underutilization of generative models by embedding a diffusion-based, logit-centric augmentation into discriminative learning through the Bernoulli-Gaussian Decision Block (BGDB). By modeling Bernoulli trial probabilities with Improved Denoising Diffusion Probabilistic Models (IDDPM) and applying the Central Limit Theorem, a single training run yields logits akin to multiple experiments, providing a Gaussian supervision signal via mean and variance targets. The authors formalize this approach with a Bernoulli approximation and a joint BGDB loss that combines task-specific objectives with diffusion-based regularization, enabling more stable and accurate classification and segmentation. Empirical results across Cityscapes, Pascal VOC, and ISIC demonstrate modest to notable improvements in segmentation (mIoU, Dice) and classification (accuracy, AUC), while ablations underscore the importance of combining loss terms and tuning hyperparameters. The work highlights BGDB’s potential to harness generative priors for discriminative tasks, balanced against practical concerns like training speed and diffusion-step constraints.

Abstract

Generative models can enhance discriminative classifiers by constructing complex feature spaces, thereby improving performance on intricate datasets. Conventional methods typically augment datasets with more detailed feature representations or increase dimensionality to make nonlinear data linearly separable. Utilizing a generative model solely for feature space processing falls short of unlocking its full potential within a classifier and typically lacks a solid theoretical foundation. We base our approach on a novel hypothesis: the probability information (logit) derived from a single model training can be used to generate the equivalent of multiple training sessions. Leveraging the central limit theorem, this synthesized probability information is anticipated to converge toward the true probability more accurately. To achieve this goal, we propose the Bernoulli-Gaussian Decision Block (BGDB), a novel module inspired by the Central Limit Theorem and the concept that the mean of multiple Bernoulli trials approximates the probability of success in a single trial. Specifically, we utilize Improved Denoising Diffusion Probabilistic Models (IDDPM) to model the probability of Bernoulli Trials. Our approach shifts the focus from reconstructing features to reconstructing logits, transforming the logit from a single iteration into logits analogous to those from multiple experiments. We provide the theoretical foundations of our approach through mathematical analysis and validate its effectiveness through experimental evaluation using various datasets for multiple imaging tasks, including both classification and segmentation.

Figures (2)

• Figure 1: Workflow for performing segmentation tasks. The total loss in the training pipeline includes $L_y$ (task-specific loss), $L_{\text{hybrid}}$ (for IDDPM), and $L_\mu + L_\sigma$ (for Bernoulli approximation). After training, only the backbones are retained for inference.
• Figure 2: Impact of hyperparameters $\lambda_2$ and $\lambda_3$ on segmentation model performance, with the other hyperparameters fixed at 1.