Research map / Probabilistic and scientific modeling / Differential equations and PDEs
PDE+: Enhancing Generalization via PDE with Adaptive Distributional Diffusion
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PDE+, which models a ResNet as a transport-equation solution and adds a learned adaptive distributional diffusion (Gaussian noise with learned scale, trained to cover AugMix-augmented samples) to improve out-of-distribution generalization on CIFAR and PACS.
Problem
Neural networks generalize poorly to distributions that differ from training data (common corruptions, domain shift); data-driven fixes such as augmentation, adversarial training, and noise injection are limited, and fixed-scale noise injection overfits to a particular noise magnitude.
Approach
Model a neural network as the solution of a transport equation and add a learnable, per-sample diffusion term (adaptive distributional diffusion) whose Gaussian scale is trained so that the diffused representation of each sample covers the representations of its semantically similar augmented samples, implemented as an ADD block after each residual block.
Key results
- On CIFAR-10-C at severity 5, PDE+ reaches 82.81% accuracy vs 76.06% for AugMix (ResNet-18).
- On CIFAR-10-C across all severities, PDE+ reaches 89.11% accuracy and 48.07 mCE, versus 86.24% and 60.44 for AugMix.
- On PACS single-source domain generalization, PDE+ averages 30.57 (photo source), 46.10 (art), 40.57 (cartoon), 35.47 (sketch), above ERM and AugMix in every split.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use PDE+ with adaptive distributional diffusion instead of RLAT AugMix | ResNet-18 · CIFAR-10-C corruption severity all | accuracy | Helped |
| Use PDE+ with adaptive distributional diffusion instead of RLAT AugMix | ResNet-18 · CIFAR-10-C corruption severity all | mCE | Helped |
| Use PDE+ with adaptive distributional diffusion instead of RLAT AugMix | ResNet-18 · CIFAR-100-C corruption severity all | accuracy: 62.77 -> 65.62 | Helped |
| Use PDE+ with adaptive distributional diffusion instead of RLAT AugMix | ResNet-18 · CIFAR-100-C corruption severity all | mCE: 77.38 -> 69.68 | Helped |
| Use PDE+ with adaptive distributional diffusion instead of RLAT AugMix | ResNet-18 · Tiny-ImageNet-C corruption severity all | accuracy: 37.92 -> 39.41 | Helped |
| Use PDE+ with adaptive distributional diffusion instead of RLAT AugMix | ResNet-18 · Tiny-ImageNet-C corruption severity all | mCE: 83.69 -> 81.80 | Helped |
| Use PDE+ instead of AugMix | ResNet-18 · PACS Photo source, target-domain average | accuracy: 26.02 -> 30.57 | Helped |
| Use PDE+ instead of AugMix | ResNet-18 · PACS Art source, target-domain average | accuracy: 43.40 -> 46.10 | Helped |
| Use PDE+ instead of AugMix | ResNet-18 · PACS Cartoon source, target-domain average | accuracy: 37.96 -> 40.57 | Helped |
| Use PDE+ instead of AugMix | ResNet-18 · PACS Sketch source, target-domain average | accuracy: 32.72 -> 35.47 | Helped |
| Use adaptive learnable diffusion instead of fixed-scale diffusion (ablation) | PDE+ · CIFAR-10-C corruption distributions | accuracy | Helped |
Surprising findings
- The paper reports PDE+ improves clean CIFAR-10 accuracy over ERM (95.59 vs 95.35), which it presents as evidence that diffusion does not trade off in-distribution performance.
Limitations
- Theoretical guarantees rely on bounded and Lipschitz assumptions and a 1D proof sketch extended to general d.
- Semantically similar samples depend on the chosen augmentation (AugMix), and overlap between augmentations and test corruptions is avoided only for most experiments.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 Tiny-ImageNet PACS
Benchmarks: CIFAR-10-C CIFAR-100-C Tiny-ImageNet-C PACS
Metrics: accuracy mCE
Baselines: ERM GradReg EnResNet RSE NFM Gaussian augmentation Mixup DeepAug AutoAug AugMix PGD RLAT
Builds on: AugMix Feynman-Kac formula Euler-Maruyama method ResNet Gaussian mixture reparameterization
Closest papers by meaning
- Boosting Generalization in Parametric PDE Neural Solvers through Adaptive Conditioning (NeurIPS 2024 · similarity 0.89)
- Adversarial Adaptive Sampling: Unify PINN and Optimal Transport for the Approximation of PDEs (ICLR 2024 · similarity 0.88)
- Generalization in VAE and Diffusion Models: A Unified Information-Theoretic Analysis (ICLR 2025 · similarity 0.88)
- DiffusionPDE: Generative PDE-Solving under Partial Observation (NeurIPS 2024 · similarity 0.88)
- Text2PDE: Latent Diffusion Models for Accessible Physics Simulation (ICLR 2025 · similarity 0.88)
- Generalization through variance: how noise shapes inductive biases in diffusion models (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.