Your ViT is Secretly a Hybrid Discriminative-Generative Diffusion Model

1 Introduction

• GenViT replaces the DDPM UNet with a single ViT for image generation, while HybViT extends this design to joint classification and generation.The paper presents these as direct connections between DDPM and ViT architectures.Source: §1 · pp. 1–2
• Reconstruction of image patches and classification operate as independent routines that train a shared set of features together.Source: §1 · p. 2
• The proposed models aim to combine DDPM-like simplicity and stability, shared statistical strengths, and ViT-derived computational efficiency and scalability.The paper contrasts these properties with collapse-prone GANs and EBMs and with the expensive sampling required by JEM.Source: §1 · p. 2
• HybViT considerably outperforms previous state-of-the-art hybrid models on both classification and generation tasks while optimizing more effectively than MCMC-based JEM/JEM++ models.The authors also report stable training without expensive MCMC sampling.Source: §1 · p. 2

2.1 Denoising Diffusion Probabilistic Models

• DDPM learns to reverse a sequence of Gaussian-noising steps, starting from noise and producing samples from the data distribution.The forward process adds Gaussian noise over timesteps, while the reverse process uses a neural network to recover data samples.Source: §2.1 · p. 2
• The DDPM variational lower bound decomposes into timestep-specific terms, with L0 modeled by a discrete decoder and LT constant and ignorable.Source: §2.1 · p. 3
• DDPM can sample an arbitrary forward timestep directly from x0 in closed form, allowing uniformly sampled timesteps to estimate the variational objective efficiently.Source: §2.1 · p. 3
• The reverse-process mean can be parameterized directly or through predicted noise, with prior work finding noise prediction effective under a reweighted loss.Source: §2.1 · p. 3
• The paper builds primarily on DDPM but notes that its objective can be replaced straightforwardly with a score-matching method.Source: §2.1 · p. 3

3 Vision Transformers

• Self-attention lets ViTs capture interactions across inputs regardless of spatial relationships, supporting non-local contextual modeling.Queries select values through keys, with attention scaled by the key dimension d_k.Source: §3 · pp. 3–4
• The GenViT and HybViT backbone feeds noisy x_t with time embedding t for generation, while HybViT computes classification logits from CLS using x0.Source: Figure 2 · p. 4
• ViT flattens an image into a sequence of patches, adds a classification token and positional embeddings, and processes the sequence with alternating attention and MLP blocks.The patch sequence length is L = H×W/P^2, and each patch has dimension P×P×C.Source: §3 · p. 4
• ViTs have been used across discriminative and generative tasks, motivating experiments with a standard ViT for generative modeling with minimal modifications.Source: §3 · p. 4
• Hybrid models jointly model density and classification with shared features, whereas this work focuses on a single ViT without an auxiliary model.The related JEM approach uses one network but relies on SGLD sampling, which remains impractical for high-dimensional data.Source: §3.1 · pp. 4–5, §3.2 · p. 5

GenViT architecture

• GenViT substitutes the DDPM UNet backbone with a single ViT and injects the timestep embedding into hidden features at every layer before reconstructing an image.The timestep embedding is produced by an MLP and modulates each layer’s hidden features.Source: §4.1 · p. 5

GenViT input and backbone

• Given DDPM input x_t, GenViT raster-scans image patches into a sequence before processing them with the ViT.Source: §4.1 · p. 5

Hybrid ViT

• HybViT replaces JEM’s energy-model likelihood optimization with GenViT’s variational objective, training classification with cross-entropy and image modeling with the diffusion simple loss.This avoids the contrastive-divergence and MCMC-based optimization used for the generative term in JEM.Source: §4.2 · p. 5
• The model extends GenViT for hybrid modeling while retaining a single ViT-based architecture for joint discriminative and generative learning.Source: §4.2 · p. 5, §4.1 · p. 5

Training objective

• Using α = 100 improves generation quality while retaining comparable classification accuracy.The resulting training pipeline is presented in Figure 3.Source: §4.2 · p. 6, Figure 3 · p. 6

5.1 Hybrid Modeling

• 95.9% accuracy and 26.4 FID on CIFAR10 let HybViT outperform JEM and JEM++ while maintaining greater training stability.The original ViT reaches comparable accuracy to WRN 28-10, and JEM-based models can diverge early.Source: §5.1 · pp. 6–7, Table 1 · p. 7

5.3 Ablation Study

• Generation quality remains below UNet-based DDPM for complex, high-resolution data, and smaller patches improve accuracy without resolving the high-resolution generation plateau.The authors identify linear and reconstruction projections, rather than multi-head self-attention alone, as possible bottlenecks.Source: §5.1 · p. 7, §5.3 · p. 10, Table 7 · p. 11

5.2.2 Out-of-Distribution Detection

• ViTs distinguish in-distribution from out-of-distribution samples better using input density scores, while HybViT and ViT have worse ECE than WRN despite appearing well-calibrated.OOD detection uses AUROC, whereas calibration uses Expected Calibration Error on CIFAR10.Source: §5.2.1 · p. 8, §5.2.2 · p. 9, Table 3 · p. 9

5.2.3 Robustness

• ViT and HybViT outperform WRN-based classifiers under white-box PGD attacks on CIFAR10.The comparison covers both L∞ and L2 perturbation constraints.Source: §5.2.3 · p. 9, Figure 7 · p. 9

5.2 Additional Analysis

• HybViT and GenViT achieve comparable CIFAR10 NLL, but both are worse than other methods in bits per dimension.The reported likelihood comparison uses test negative log-likelihood.Source: §5.2.4 · p. 9, Table 4 · p. 10

5.3 Ablation Study

• Training from 100 to 300 epochs substantially improves classification and generation, while larger α and model capacity favor generation with only small accuracy costs.More epochs slightly improve generation beyond 300, and larger depth or hidden dimensions improve both accuracy and generation quality.Source: §5.3 · p. 10, Table 5 · p. 10, Table 6 · p. 11

• Generation quality is relatively low compared with pure-generation state-of-the-art models.Source: §6 · p. 11
• HybViT requires more training iterations to reach high classification performance than pure classification models.Source: §6 · p. 11
• Inference sampling is slow, typically requiring T ≥1000, whereas GAN requires one forward pass.Source: §6 · p. 11
• The authors hope these limitations motivate improvements in sampling speed and generative quality.Source: §6 · p. 11

• GenViT integrates a single ViT into DDPM, while HybViT trains a hybrid discriminative-generative model.Source: §7 · p. 11
• The models achieve state-of-the-art results in most image-classification and image-generation tasks across multiple benchmark datasets.Source: §7 · p. 11
• HybViT provides stable training and outperforms previous state-of-the-art hybrid models on discriminative and generation tasks.Source: §7 · p. 11
• Runtime comparisons on CIFAR10 use 1170 iterations per ViT epoch and 390 iterations per WRN-based-model epoch.Source: Table 8 · p. 12
• VAEBM requires an additional 25,000 CIFAR-10 iterations, or 16 epochs, after NVAE pretraining on one GPU.Source: §7 · p. 12

• CIFAR10 contains 60,000 32×32 RGB images across 10 classes, split into 50,000 training and 10,000 test images.Source: A Image Datasets · p. 15
• CIFAR100 contains 60,000 32×32 RGB images across 100 classes, with 500 training and 100 test images per class.Source: A Image Datasets · p. 15
• STL10 provides 500 training images and 800 test images per class across 10 classes.Source: A Image Datasets · p. 15
• Tiny-ImageNet contains 100,000 downsized 64×64 color images across 200 classes, with training, validation, and test splits.Source: A Image Datasets · p. 15
• CelebA-HQ is represented by a downsampled 128×128 human-face dataset.Source: A Image Datasets · p. 15
• ImageNet 32×32 is a 1,000-class downsampled ImageNet variant with the original dataset's image count.Source: A Image Datasets · p. 15

• Experiments follow vanilla ViT and DDPM settings, use SGD with initial learning rate 0.1, and apply cosine learning-rate decay.Source: B Experimental Details · p. 15
• Table 9 lists the hyper-parameters used for ViT, GenViT, and HybViT.Source: Table 9 · p. 16
• Figure 8 tracks HybViT classification accuracy and FID across training epochs on CIFAR10 and ImageNet 32x32.Source: Figure 8 · p. 16
• Figure 9 compares samples with FID=40 and FID=20, whose difference is visually imperceptible to humans.Source: Figure 9 · p. 16

C.1 Qualitative Analysis of Samples

• Generation-quality improvement becomes relatively small after 10% of training, with FID=40 and FID=20 samples nearly visually indistinguishable.The authors identify faster convergence as an important direction for future work.Source: §C.1 · p. 17
• CIFAR10 qualitative analysis includes class-conditional samples and class-wise histograms, with some worst Plane samples becoming completely blank.Source: §C.1 · p. 17, Figure 10 · p. 17, Figure 11 · p. 17

Data Augmentation

• Strong augmentation is essential for high classification performance, while its effect on generation is negative but tiny.Source: §C.2 · p. 18

OOD Detection and Runtime

• The analysis also evaluates OOD detection with classifier maximum probability and reports memory-efficient, faster training on a single GPU.OOD results are reported as AUROC, and runtime is measured on ImageNet 32x32 for 100 epochs.Source: §C.3 · p. 18, Table 11 · p. 18, §C.6 · p. 19

Robustness

• HybViT with weak augmentation is much more robust than other models, especially under L2 attacks.The authors hypothesize that diffusion exposes HybViT to noisy samples, implicitly improving flatness and robustness.Source: §C.4 · p. 18, Table 12 · p. 18

Calibration

• Strong augmentation better calibrates ViT and HybViT predictions but makes them under-confident relative to the calibration comparison.The comparison includes WRN, JEM, and corresponding ViTs trained without strong augmentation.Source: §C.5 · p. 19, Figure 12 · p. 19

Additional Generated Samples

• Additional samples are shown for CIFAR10, CIFAR100, ImageNet 32x32, TinyImageNet, STL10, and CelebA 128.Further examples are provided for ImageNet 128x128 and vanilla ImageNet 224x224.Source: D Additional Generated Samples · p. 19, Figure 13 · p. 20
• Preliminary ImageNet 128x128 and vanilla ImageNet 224x224 generation has very low quality under limited computation resources.The reported patch sizes are 8 for ImageNet 128 and 14 for ImageNet 224.Source: D Additional Generated Samples · p. 19, Figure 14 · p. 20

Class-Conditional Samples

• HybViT class-conditional samples are presented separately for Plane, Car, Bird, Cat, Deer, Dog, Frog, Horse, Ship, and Truck.Source: Figure 15 · p. 20, Figure 16 · p. 21, Figure 17 · p. 21
• The supplementary figures extend class-conditional visualization across all ten CIFAR10 classes, including Frog, Horse, Ship, and Truck.Source: Figure 21 · p. 22, Figure 22 · p. 22, Figure 23 · p. 22