State-Based Diffusion Models

• State-based diffusion models are defined as frameworks where a forward corruption process and learned reverse denoising operate on explicit, structured state spaces.
• They integrate diverse approaches—from categorical token diffusion to state-space backbones in SSMs—to improve long-range coherence, transition expressivity, and memory efficiency.
• Applications span image and video generation, reinforcement learning, and scientific modeling, consistently achieving competitive results with reduced computational complexity.

Searching arXiv for papers on state-based diffusion models and closely related formulations. Search query: "state-based diffusion models diffusion state space models arXiv" State-based diffusion models are diffusion-based systems in which the relevant object is explicitly a state: a discrete Markov state, a latent dynamical state, a successor-state distribution, an intermediate diffusion state, a hidden state in a structured state space model, or a non-Euclidean scientific state such as a quantum pure state. The label is therefore not a single architectural prescription but a family resemblance across several research lines. What unifies these lines is a forward corruption or diffusion process defined on a state space and a learned reverse or denoising process; what varies is whether “state” denotes the domain being modeled, the internal backbone carrying memory, or the dynamical variable whose transitions are learned (Pauline et al., 4 Dec 2025).

1. Terminological scope and internal taxonomy

A common source of confusion is that the phrase state-based diffusion model is used in several distinct but compatible senses. In discrete generative modeling, it refers to diffusion on categorical state spaces. In sequence modeling, it often refers to diffusion denoisers built on structured state space models (SSMs) such as S4 or Mamba. In reinforcement learning, it can mean diffusion over future-state distributions or optimization at intermediate diffusion states. In geometric and scientific settings, it can mean diffusion directly on a manifold-valued state space such as $\mathbb{CP}^{d-1}$ (Austin et al., 2021, Fei et al., 2024, Schramm et al., 2024, Xu et al., 5 May 2026).

Usage of “state-based”	State object	Representative papers
Discrete-state diffusion	Categorical tokens or symbols	(Austin et al., 2021, Pauline et al., 4 Dec 2025)
State-space backbone	Hidden states in SSM denoisers	(Fei et al., 2024, Alcaraz et al., 2022, Oshima et al., 2024, Yuan et al., 27 Jul 2025)
Latent dynamics modeling	Latent sequential states with learned transitions	(Savov et al., 28 May 2025, Ruder et al., 19 Jun 2026)
RL state modeling	Successor states, compositional environment states, or intermediate diffusion states	(Schramm et al., 2024, Duan et al., 22 Jan 2025, Oba et al., 6 Feb 2026)
Manifold or scientific state diffusion	Pure quantum states or state frequencies	(Xu et al., 5 May 2026, Soewongsono et al., 2024)

This multiplicity is not merely terminological. It reflects different answers to three questions: what is being diffused, where long-range structure is stored, and what role denoising plays. Some papers diffuse the state itself, as in discrete D3PMs or pure-state manifold diffusion. Others keep a standard diffusion objective but replace the denoiser with an SSM backbone. Still others use diffusion to parameterize state transitions inside a broader state-space model or world model (Pauline et al., 4 Dec 2025, Ruder et al., 19 Jun 2026).

2. General-state-space formulation

The most general formulation treats diffusion as a Markov process on a state space $X^d$. In discrete time, the forward process is

$$q(x_{0:T}) = q(x_0)\prod_{t=1}^T q(x_t\mid x_{t-1}),$$

and the learned reverse process is

$$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t).$$

This template covers continuous Gaussian corruption, categorical transition kernels, continuous-time SDE limits in $\mathbb{R}^d$, and continuous-time Markov chains on finite alphabets. The corresponding marginal-evolution equations are the Fokker–Planck equation in continuous spaces and the master equation in discrete spaces, unified by the generator identity $\partial_t q_t = (L_t)^* q_t$ (Pauline et al., 4 Dec 2025).

Within discrete state spaces, D3PMs made the forward kernel itself a modeling choice rather than a fixed multinomial replacement rule. For a categorical variable with $K$ states, the one-step corruption is

$$q(x_t\mid x_{t-1})=\mathrm{Cat}(x_t;\;p=x_{t-1}Q_t),$$

with cumulative kernel $\overline Q_t = Q_1Q_2\cdots Q_t$. The literature studied uniform transitions, Gaussian-like kernels for ordinal structure, nearest-neighbor transitions in embedding space, and absorbing-state corruption. Absorbing-state diffusion is especially significant because it connects discrete diffusion to masked generative modeling, while the hybrid objective

$$L_\lambda = L_{\mathrm{vb}} + \lambda\, \mathbb E[-\log \widetilde p_\theta(x_0\mid x_t)]$$

showed that the auxiliary cross-entropy term can stabilize training and improve sample quality (Austin et al., 2021).

A central implication of this general-state-space view is that diffusion is not inherently tied to Euclidean images. The same variational logic yields denoising score matching in continuous spaces and denoising score entropy or masked-language-model reductions in discrete spaces. This suggests that “state-based diffusion” is best understood as a probabilistic framework whose specific form is determined by the forward corruption process and the geometry or algebra of the underlying state space, rather than by any single network family (Pauline et al., 4 Dec 2025).

3. State-space backbones inside diffusion denoisers

A major strand of the literature uses state space in an architectural sense: the reverse denoiser is built from SSM blocks rather than from a conventional CNN U-Net or Transformer U-Net. In image generation, Diffusion State Space Models (DiS) replace the usual backbone with stacked bidirectional Mamba-style blocks. DiS treats noisy image patches, the timestep, and the conditioning signal as tokens, appending timestep and condition as two supplementary tokens to the image-token sequence. The architecture consists of patchify and linear embedding, a stack of SSM blocks, a long skip connection implemented by concatenation plus projection, and a linear decoder. On ImageNet in latent space, DiS-H/2 reports FID 2.10 at $X^d$0 and FID 2.88 at $X^d$1, while the paper’s scalability analysis emphasizes that higher Gflops correlate with lower FID and that the SSM complexity is linear in token count $X^d$2, in contrast to the quadratic term in self-attention (Fei et al., 2024).

In video generation, the same idea appears in a more localized form. Rather than replacing the whole backbone, the temporal attention layers in a video diffusion U-Net are replaced by a bidirectional S4D-based temporal SSM block: $X^d$3 The forward and backward streams are summed, then passed through an MLP. This design is not a naïve drop-in replacement: the paper shows that bidirectionality is essential and that removing the MLP hurts performance because the SSM alone has SISO-style limitations in cross-channel mixing. For sequences up to 256 frames, SSM-based models require less memory to achieve the same FVD as attention-based models; on MineRL at 400 frames, attention could not be trained due to memory limits, while the SSM model achieved FVD 972.306 (Oshima et al., 2024).

Time-series imputation and forecasting supply an earlier and highly explicit version of this architectural move. SSSD replaces DiffWave-style dilated convolutions with S4 layers, conditions on observed values and masks, and distinguishes a $X^d$4 regime that diffuses the whole signal from a preferred $X^d$5 regime that diffuses only the missing region. The paper reports especially strong results on long contiguous gaps such as blackout missingness and forecasting, where state-space memory is most useful; on PTB-XL with 20% blackout missingness, SSSD reports MAE 0.0324 and RMSE 0.0832 versus CSDI MAE 0.1054 and RMSE 0.2254 (Alcaraz et al., 2022).

In symbolic music, SMDIM combines a discrete D3PM with a hierarchical Mamba backbone and a Mamba-FeedForward-Attention block. The order $X^d$6 is reported to work better than attention-first orderings, reflecting a division of labor between global context, nonlinear refinement, and local token precision. On MAESTRO, POP909, and FolkDB, SMDIM reports average OA values of 0.850, 0.885, and 0.854 respectively, while also reducing step time and GPU memory relative to SCHmUBERT at sequence length 2048 (Yuan et al., 27 Jul 2025).

4. Long-context memory and latent state dynamics

Another interpretation of state-based diffusion places the state not in the denoiser backbone but in an explicit latent dynamical memory. StateSpaceDiffuser is exemplary: it combines a Long-Context Branch implemented with Mamba and a Generative Branch built on DIAMOND, a UNet-based EDM diffusion model. The long-context branch processes full-frame history together with actions, using Cosmos tokenizer features and a 16-dimensional action embedding; the diffusion branch sees only a short context of four frames plus actions, predicts the next frame at low resolution, and uses a second diffusion model for upsampling. Fusion occurs by taking the last four SSM output features, normalizing them, passing them through a two-layer MLP with SiLU, separately processing action embeddings with a similar MLP, and concatenating the results into the conditioning vector (Savov et al., 28 May 2025).

This architecture directly addresses the failure mode that diffusion-only world models lose track of long-term context because they condition on only a short recent window. The evaluation protocol is correspondingly memory-centric: in MiniGrid, the model is tested on 100-step sequences consisting of 50 forward and 50 backward steps, and the reverse trajectory must reconstruct previously seen content. With context length 16, StateSpaceDiffuser reports Avg PSNR 41.01, Final PSNR 40.55, and SSIM 0.98, compared with DIAMOND’s Avg PSNR 27.13, Final PSNR 25.44, and SSIM 0.95; the headline claim is an 8.9 dB average PSNR improvement over the diffusion baseline (Savov et al., 28 May 2025).

A more structural treatment of latent dynamics appears in DDSSM, which replaces the Gaussian transition distribution of a deep state-space model with a conditional diffusion model. Starting from the factorization

$X^d$7

DDSSM defines the transition $X^d$8 by marginalizing a reverse diffusion chain over noisy latent states. The model is trained end-to-end under an ELBO whose negative form decomposes into initialization, reconstruction, diffusion, and posterior-entropy terms. The paper’s main methodological claim is that this resolves the open problem of jointly training an autoencoder and a diffusion model on sequential data, rather than relying on a two-stage latent diffusion pipeline. Empirically, on a synthetic bimodal-noise random walk, DDSSM captures the bimodal one-step predictive distribution and achieves lower JSD than a DKF baseline, though the paper also notes that the evidence is limited to synthetic validation and that the emission model remains Gaussian (Ruder et al., 19 Jun 2026).

These two lines suggest complementary uses of state. StateSpaceDiffuser externalizes state as long-context memory that conditions a high-fidelity diffusion renderer, whereas DDSSM internalizes diffusion as the transition law of the latent state itself. A plausible implication is that “state-based diffusion” can either augment diffusion with persistent memory or absorb diffusion into the dynamical model, depending on whether the main bottleneck is long-context conditioning or transition expressivity.

5. Reinforcement learning and sequential decision-making

In reinforcement learning, state-based diffusion models often depart from the now-common use of diffusion as a direct action generator. Bellman Diffusion Models instead treat the successor state measure of a policy as the object learned by diffusion: $X^d$9 The key step is the Bellman flow constraint

$$q(x_{0:T}) = q(x_0)\prod_{t=1}^T q(x_t\mid x_{t-1}),$$0

which induces a Bellman-style bootstrapped denoising target. The resulting diffusion loss has one term for immediate next-state denoising and one term for a target-network bootstrap, making the method a diffusion analogue of temporal-difference learning over future-state distributions rather than actions (Schramm et al., 2024).

A different state-based interpretation arises in zero-shot generalization to unseen combinations of known elements. The work on state combinatorial generalization formalizes out-of-combination generalization as a support mismatch, $$q(x_{0:T}) = q(x_0)\prod_{t=1}^T q(x_t\mid x_{t-1}),$$1, while assuming that all base elements have been seen during training. The argument is that value-based RL fails because value estimates become unreliable in unsupported states, whereas a conditional diffusion model trained by behavior cloning on expert trajectories can condition on compositional information and generate plausible trajectories in those states. On HighwayEnv Roundabout, the paper reports OOD crashes dropping from 36.0 for PPO to 20.0 for diffusion, with OOD reward increasing from 8.19 to 8.81 (Duan et al., 22 Jan 2025).

The notion of state becomes even more literal in DiSPO for masked diffusion LLMs. Here, each partially masked sequence $$q(x_{0:T}) = q(x_0)\prod_{t=1}^T q(x_t\mid x_{t-1}),$$2 is treated as a reinforcement-learning state, and the joint filling of currently masked positions is treated as an action: $$q(x_{0:T}) = q(x_0)\prod_{t=1}^T q(x_t\mid x_{t-1}),$$3 DiSPO branches from rollout-cached logits at selected intermediate states, scores the resulting completions with the same terminal reward, and updates only the newly filled tokens. The paper derives a fixed-state objective $$q(x_{0:T}) = q(x_0)\prod_{t=1}^T q(x_t\mid x_{t-1}),$$4 and a policy-gradient estimator whose expected gradient aligns with that objective; it reports consistent improvements over diffu-GRPO on GSM8K, MATH500, 4×4 Sudoku, and Countdown under matched rollout compute and optimizer steps (Oba et al., 6 Feb 2026).

Across these papers, the state-based perspective in decision-making has three distinct meanings: the diffusion target can be a state distribution, the conditioning variable can encode combinatorial state structure, or the intermediate denoising configuration can itself be optimized as a state. This suggests that RL has become a particularly fertile domain for expanding the semantics of “state-based diffusion.”

6. Manifold-valued and scientific state spaces

The most geometrically explicit instance of state-based diffusion is SSDM for quantum pure-state ensemble generation. Here the state space is not $$q(x_{0:T}) = q(x_0)\prod_{t=1}^T q(x_t\mid x_{t-1}),$$5 but the complex projective manifold

$$q(x_{0:T}) = q(x_0)\prod_{t=1}^T q(x_t\mid x_{t-1}),$$6

equipped with the Fubini–Study metric and geodesic distance $$q(x_{0:T}) = q(x_0)\prod_{t=1}^T q(x_t\mid x_{t-1}),$$7. The forward process is an intrinsic manifold diffusion, optionally realized through a stochastic Schrödinger equation on the Hilbert sphere and then quotiented by global phase. Reverse-time sampling uses the Riemannian score,

$$q(x_{0:T}) = q(x_0)\prod_{t=1}^T q(x_t\mid x_{t-1}),$$8

and training avoids intractable transition densities by introducing a local-time objective in FS normal coordinates with an analytic teacher score derived from a local Euclidean Ornstein–Uhlenbeck approximation (Xu et al., 5 May 2026).

This paper is significant because it shows that respecting state-space geometry can be central rather than incidental. The empirical study, conducted for $$q(x_{0:T}) = q(x_0)\prod_{t=1}^T q(x_t\mid x_{t-1}),$$9 qubits, evaluates average state fidelity, observable moment mismatch, kernel MMD with overlap kernel, and entanglement statistics. SSDMs are reported to be consistently the best or among the best across all metrics, with advantages increasing as system size grows, and SSDM-generated quantum feature states improve downstream kernel alignment, kernel gap, mean margin, and classifier generalization under limited data (Xu et al., 5 May 2026).

A different scientific use of state-based diffusion appears in phylogenetic modeling of state-dependent speciation and extinction. Here “diffusion” does not denote denoising generative modeling but a diffusion approximation to the forward-in-time evolution of state counts $$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t).$$0 and normalized state frequencies $$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t).$$1 under ClaSSE or GeoSSE dynamics: $$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t).$$2 The framework is used to compare tree-based and diffusion-based simulation, compute stationary state frequencies, infer compatible rate parameters, and estimate time to stationarity. The paper reports that under large initial clade sizes the end-of-run means from diffusion-based and tree-based simulations are statistically indistinguishable by Welch’s t-test, while variance ratios need not equal 1 (Soewongsono et al., 2024).

This juxtaposition is instructive. In one case, diffusion is a generative process on a curved manifold of quantum states; in the other, it is an SDE approximation for macroevolutionary state frequencies. The two uses are mathematically related through stochastic-process language, but they belong to different methodological traditions. A common misconception is therefore to treat all “state-based diffusion” papers as members of a single machine-learning architecture family; the scientific literature shows that the phrase also names a broader class of diffusion processes defined on explicit state variables.

7. Empirical regularities, misconceptions, and limitations

Several empirical regularities recur across otherwise different state-based diffusion models. First, replacing temporal or sequence-wide attention with SSM components is repeatedly motivated by linear memory or near-linear complexity. In video diffusion, the contrast is explicit: attention has memory $$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t).$$3 while SSMs have memory $$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t).$$4 and remain parallelizable; in symbolic music, most of the computation is shifted from $$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t).$$5 self-attention to Mamba-dominated processing; in image generation, DiS derives a linear-in-token SSM term versus a quadratic self-attention term (Oshima et al., 2024, Yuan et al., 27 Jul 2025, Fei et al., 2024).

Second, the literature does not support the stronger claim that attention is uniformly obsolete. The video paper states that its point is not that SSM always beats attention in short-video FVD, but that it delivers competitive quality while enabling much longer generation, and it explicitly does not claim that attention is obsolete. The music paper retains a self-attention stage inside each MFA block for local precision. DGSSM, although framed as a diffusion-guided state space framework for multimodal salient object detection, uses diffusion as structural guidance and iterative refinement rather than as a full standalone generative sampler; on DUTS it reports $$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t).$$6, $$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t).$$7, and $$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t).$$8 with 38.92M parameters, outperforming Samba while remaining more compact (Oshima et al., 2024, Yuan et al., 27 Jul 2025, Ghosh et al., 19 Apr 2026).

Third, hybridization is often essential. StateSpaceDiffuser shows that diffusion-only baselines have high local fidelity but poor long-horizon consistency, whereas state-space-only models are blurrier or artifact-prone; the best results come from fusing long-context state memory with a diffusion renderer. In video diffusion, the authors report that a naïve drop-in SSM replacement performs poorly and that bidirectionality plus a post-SSM MLP are crucial. In DDSSM, joint end-to-end training is presented as preferable to two-stage latent diffusion because otherwise the latent space may be over-regularized (Savov et al., 28 May 2025, Oshima et al., 2024, Ruder et al., 19 Jun 2026).

Limitations are equally consistent. Some models are validated only on unconditional generation or synthetic data; the video SSM paper focuses on unconditional video generation, and DDSSM reports only a synthetic bimodal sequence study. Some require specialized pretraining or multistage training, as in StateSpaceDiffuser. Some remain computationally heavy at sampling time because diffusion steps are still expensive even when the state representation is improved. In decision-making, the combinatorial-generalization work explicitly does not address unseen base objects, only unseen combinations of seen ones, and DiSPO notes non-negligible implementation overhead despite matched rollout counts (Oshima et al., 2024, Ruder et al., 19 Jun 2026, Savov et al., 28 May 2025, Duan et al., 22 Jan 2025, Oba et al., 6 Feb 2026).

Taken together, these results indicate that state-based diffusion models are best viewed not as a single canonical model class but as a design pattern. The design pattern has three recurring motives: impose a structured state space on the data or latent variables, use state evolution to improve long-range coherence or transition expressivity, and preserve diffusion’s strength as a flexible denoising or generative mechanism. This suggests why the label has spread across image, video, time series, symbolic music, world modeling, reinforcement learning, quantum representation learning, dense prediction, and evolutionary dynamics without collapsing into a single formal definition.