Steering Diffusion Models via Class-Contrastive Influence for Few-Shot Medical Classification

Technical Summary: Steering Diffusion Models via Class-Contrastive Influence for Few-Shot Medical Classification

Problem Statement
In few-shot medical image classification, labeled data is scarce, making data augmentation essential. While off-the-shelf diffusion models can generate synthetic training samples, existing approaches primarily focus on improving the realism, diversity, or domain adaptation of these images. The authors argue that these metrics do not guarantee utility for the downstream classification task. A significant discrepancy exists where generated images of comparable visual quality yield vastly different classification performance. The core problem identified is the lack of a principled criterion to measure and optimize the "task usefulness" of synthetic samples. The paper posits that the central limitation of current practices is not generation quality alone, but the failure to generate samples that specifically aid the classifier in refining its decision boundary.

Methodology
The proposed solution, Class-Contrastive Influence (C2I), shifts the optimization objective from image quality to the gradient-based influence of synthetic samples on the classifier.

1. Class-Contrastive Influence (C2I) Criterion:

   • The method leverages gradient-based influence estimation to quantify how a training sample affects validation loss.
   • The authors observe that raw, class-agnostic influence scores do not correlate with classification performance because validation sets contain conflicting signals from different classes.
   • Instead, they define C2I based on the influence gap: a sample is deemed useful if its loss gradient aligns positively with validation gradients from the same class and opposes (aligns negatively with) gradients from other classes.
   • For binary classification, C2I is formulated as the squared difference between the mean influence on the positive and negative validation sets, normalized by the variance. For multi-class tasks, it is defined as the softmax score of the mean influence on the matching class relative to all classes.

2. Theoretical Connection to Sample Hardness:

   • Theoretical analysis (Theorem 1) demonstrates that maximizing C2I in a logistic regression setting drives the generated sample toward the global mean of the validation set.
   • Since the global mean of two class clusters typically lies near the decision boundary, maximizing C2I inherently selects "hard" examples—samples that are difficult to classify and lie close to the boundary.
   • Training on these hard, boundary-proximal examples helps refine the decision boundary, improving generalization and robustness.

3. Reinforcement Learning (RL) Fine-Tuning:

   • To operationalize this, the authors propose a fine-tuning scheme for pre-trained Text-to-Image (T2I) diffusion models (specifically Stable Diffusion 2.1).
   • Reward Function: The C2I score serves as the reward signal in a Reinforcement Learning framework. The generator is fine-tuned to maximize the expected C2I reward.
   • Optimization: The authors utilize Denoising Diffusion Policy Optimization (DDPO) to update the diffusion model's parameters (specifically LoRA weights) based on the C2I reward. This steers the generation process toward the "class-informative" regions of the data manifold without requiring architectural changes to the generator or the classifier.
   • Workflow: A classifier (ViT-B/16) is first fine-tuned on the few-shot training set. Its gradients on a held-out validation set are used to compute the C2I reward for generated batches, which then guides the RL fine-tuning of the diffusion model.

Key Contributions

• Formalization of Effectiveness: The paper formalizes sample effectiveness through C2I, demonstrating that high-C2I samples correspond to hard, boundary-proximal examples rather than just realistic ones.
• RL-Based Generation Scheme: It introduces a plug-and-play RL framework that uses C2I as a reward to steer off-the-shelf diffusion models toward generating task-useful, class-informative samples.
• Empirical Validation: The method is validated across multiple few-shot medical imaging benchmarks (BreastMNIST, DermaMNIST, PneumoniaMNIST), showing consistent improvements over standard augmentation and other diffusion-based baselines.

Experimental Results
The method was evaluated on three MedMNIST benchmarks under few-shot settings (16 or 32 labeled samples per class).

• Performance: The C2I-guided generator consistently outperformed baselines, including standard transformations (RandAugment, RandomErasing, Mixup) and other diffusion-based augmentation methods (DataDream, Dataset Expansion, DistDiff).
  • On average, the proposed method achieved the highest AUC across different backbones (ViT-B/16 and ResNet-18).
  • Notably, the method improved performance on ResNet-18, indicating that the synthesized samples are broadly informative and transfer across model architectures.
• Robustness: Models trained with C2I-guided augmentation demonstrated superior robustness under various noise conditions (salt-and-pepper, Gaussian blur, JPEG compression), suggesting the generation of a more stable decision boundary.
• Hardness Correlation: Experiments confirmed that as RL fine-tuning progresses, the C2I reward increases, and the features of generated images move closer to the validation feature centroids, confirming the generation of harder examples.
• Ablation: Increasing the number of RL training epochs and the number of synthesized images generally improved downstream AUC, whereas some baselines (e.g., Dataset Expansion) showed performance degradation with increased synthetic data.

Significance and Claims
The paper claims that synthetic augmentation is most effective when guided by task usefulness rather than image quality alone. By aligning synthetic data with validation gradients, C2I transforms generic diffusion models into targeted generators that strengthen decision boundaries where labeled data is most scarce.

The authors acknowledge limitations, noting that their method introduces additional computational overhead compared to using original training data alone, primarily due to the RL fine-tuning process (requiring approximately 5 GPU hours on an A100). However, they assert that this cost is justified by the consistent gains in accuracy and robustness in low-data regimes. The study suggests a paradigm shift in using diffusion models for augmentation: optimizing generation for its downstream influence on the classifier.