ASSERT: Adaptive Stochastic Sampling for Robust Diffusion Models on Analog Compute-in-Memory Hardware

• URL: http://arxiv.org/abs/2609.00955v1
• Date: Tue, 01 Sep 2026 09:13:10 GMT
• Title: ASSERT: Adaptive Stochastic Sampling for Robust Diffusion Models on Analog Compute-in-Memory Hardware
• Abstract summary: Diffusion models achieve strong image generation quality but incur high iterative denoising costs.<n>We show that the early, high-noise denoising stage is substantially more vulnerable than the final refinement stage.<n>We propose ASSERT, a training sampler that uses higher deterministicity early and smoothly denoising.
• Score: 13.123394042691642
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Diffusion models achieve strong image generation quality but incur high iterative denoising costs. Analog compute-in-memory (CIM) can accelerate matrix-vector multiplications, yet spatial memory variations perturb weights and accumulate during sampling. Unlike conventional neural networks, diffusion models' temporal sensitivity to hardware noise remains underexplored. We investigate diffusion inference using a noise model calibrated and validated against measurements collected from multiple physical CIM chips. Our results show that the early, high-noise denoising stage is substantially more vulnerable than the final refinement stage. A first-order trajectory analysis attributes this behavior to the repeated propagation of correlated prediction errors induced by a fixed hardware mapping. Based on this observation, we propose ASSERT, a training-free sampler that uses higher stochasticity early and smoothly transitions to deterministic denoising. The injected stochasticity changes subsequent activation trajectories and thereby reduces their alignment with persistent spatial errors. Across the evaluated settings, ASSERT achieves up to 2.58$\times$ lower FID than deterministic DDIM on high-resolution datasets and 7.68$\times$ lower FID in the CIFAR-10 step-count study, without changing model parameters or the number of network evaluations.

Related papers

• SATB-VR: Training Few-Step Video Restoration Diffusion Model using SNR-Aware Trajectory Blending [55.764018145256216]
  We present SATB-VR, a few-step paradigm that jump-starts the denoising process via an auxiliary predictor.<n>We show that SATB-VR performs favorably against existing approaches on synthetic, real-world, and AIGC benchmarks.
  arXiv  Detail & Related papers  (2026-06-27T01:32:16Z)
• Precise: SDE-Consistent Stochastic Sampling for RL Post-Training of Flow-Matching Models [56.67321805551389]
  Reinforcement learning (RL) has become an effective way to improve prompt alignment and perceptual quality in diffusion and flow-matching generators.<n>The sampler, controlling the exploration behavior and denoising dynamics, is thus part of the policy.<n>We propose a new sampler that balances effective exploration with stability.
  arXiv  Detail & Related papers  (2026-05-22T11:37:22Z)
• Transformers Learn the Optimal DDPM Denoiser for Multi-Token GMMs [13.741630476895773]
  This paper provides the first convergence analysis for training transformer-based diffusion models.<n>We consider the population Denoising Diffusion Probabilistic Model (DDPM) objective for denoising data.<n>A deeper investigation reveals that trained a mean denoising module enables the trained model to approximate the oracle Minimum Mean Squared Error (MMSE) of the injected noise in the diffusion steps.
  arXiv  Detail & Related papers  (2026-04-11T07:46:15Z)
• Not All Denoising Steps Are Equal: Model Scheduling for Faster Masked Diffusion Language Models [6.4453302264198165]
  We exploit the flexibility of the diffusion framework and study model scheduling, where a smaller MDLM replaces the full model at a subset of denoising steps.<n>We show that early and late denoising steps are substantially more robust to such replacement than middle steps, enabling up to a 17% reduction in FLOPs.<n>Our results suggest that simple, architecture-agnostic scheduling rules can significantly accelerate MDLM sampling while largely preserving generation quality.
  arXiv  Detail & Related papers  (2026-02-04T13:04:58Z)
• Automated Tuning for Diffusion Inverse Problem Solvers without Generative Prior Retraining [4.511561231517167]
  Diffusion/score-based models have emerged as powerful generative priors for solving inverse problems.<n>We propose Zero-shot Adaptive Diffusion Sampling (ZADS), a test-time optimization method that tunes fidelity weights across arbitrary noise schedules.<n>ZADS consistently outperforms both traditional compressed sensing and recent diffusion-based methods.
  arXiv  Detail & Related papers  (2025-09-11T22:22:32Z)
• Noise Conditional Variational Score Distillation [60.38982038894823]
  Noise Conditional Variational Score Distillation (NCVSD) is a novel method for distilling pretrained diffusion models into generative denoisers.<n>By integrating this insight into the Variational Score Distillation framework, we enable scalable learning of generative denoisers.
  arXiv  Detail & Related papers  (2025-06-11T06:01:39Z)
• Test-Time Scaling of Diffusion Models via Noise Trajectory Search [10.8507840358202]
  We introduce an $epsilon$-greedy search algorithm that globally explores at extreme timesteps and locally exploits during the intermediate steps where de-mixing occurs.<n>Experiments on EDM and Stable Diffusion reveal state-of-the-art scores for class-conditioned/text-to-image generation.
  arXiv  Detail & Related papers  (2025-05-24T19:13:29Z)
• Spatial Reasoning with Denoising Models [49.83744014336816]
  We introduce a framework to perform reasoning over sets of continuous variables via denoising generative models.<n>For the first time, that order of generation can successfully be predicted by the denoising network itself.<n>Using these findings, we can increase the accuracy of specific reasoning tasks from 1% to >50%.
  arXiv  Detail & Related papers  (2025-02-28T14:08:30Z)
• Robust Representation Consistency Model via Contrastive Denoising [83.47584074390842]
  randomized smoothing provides theoretical guarantees for certifying robustness against adversarial perturbations.<n> diffusion models have been successfully employed for randomized smoothing to purify noise-perturbed samples.<n>We reformulate the generative modeling task along the diffusion trajectories in pixel space as a discriminative task in the latent space.
  arXiv  Detail & Related papers  (2025-01-22T18:52:06Z)
• Inference Stage Denoising for Undersampled MRI Reconstruction [13.8086726938161]
  Reconstruction of magnetic resonance imaging (MRI) data has been positively affected by deep learning. A key challenge remains: to improve generalisation to distribution shifts between the training and testing data.
  arXiv  Detail & Related papers  (2024-02-12T12:50:10Z)
• Adaptive Multi-step Refinement Network for Robust Point Cloud Registration [82.64560249066734]
  Point Cloud Registration estimates the relative rigid transformation between two point clouds of the same scene.<n>We propose an adaptive multi-step refinement network that refines the registration quality at each step by leveraging the information from the preceding step.<n>Our method achieves state-of-the-art performance on both the 3DMatch/3DLoMatch and KITTI benchmarks.
  arXiv  Detail & Related papers  (2023-12-05T18:59:41Z)
• SMRD: SURE-based Robust MRI Reconstruction with Diffusion Models [76.43625653814911]
  Diffusion models have gained popularity for accelerated MRI reconstruction due to their high sample quality. They can effectively serve as rich data priors while incorporating the forward model flexibly at inference time. We introduce SURE-based MRI Reconstruction with Diffusion models (SMRD) to enhance robustness during testing.
  arXiv  Detail & Related papers  (2023-10-03T05:05:35Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.