---
title: 'STADI: Spatio-Temporal Adaptive Diffusion Inference'
url: https://www.emergentmind.com/topics/spatio-temporal-adaptive-diffusion-inference-stadi
type: topic
---

# STADI: Spatio-Temporal Adaptive Diffusion Inference

Searching arXiv for the STADI paper and directly related antecedents.
First, the exact STADI paper.
Then, the two earlier spatio-temporal adaptive diffusion papers that are conceptually related but use different names/acronyms.
Finally, a prior patch-parallel diffusion inference line for context.
Spatio-Temporal Adaptive Diffusion Inference (STADI) is a distributed inference framework for accelerating diffusion model sampling on heterogeneous multi-GPU systems. In the exact arXiv usage of the acronym, STADI denotes a scheduling method that exploits two dimensions of diffusion inference simultaneously: the temporal dimension, defined by the number of denoising steps each GPU performs, and the spatial dimension, defined by the image patch size each GPU processes [2509.04719]. The term also sits within a broader family of spatio-temporal adaptive diffusion methods in which “spatio-temporal” refers to joint modeling or conditioning over time and space, as in Spatio-temporal Diffusion Point Processes (DSTPP) for spatio-temporal point processes and STADMs for EEG super-resolution; these are conceptually related antecedents, but they do not use the exact STADI name [2305.12403] [2407.03089].

## 1. Definition and conceptual scope

In its precise formulation, STADI addresses workload imbalance during diffusion inference in heterogeneous or load-imbalanced multi-GPU environments. The framework is motivated by the observation that diffusion sampling is inherently iterative and that existing parallel inference schemes often exploit only spatial partitioning, typically by assigning fixed image patches to different GPUs and synchronizing at every denoising step. In such settings, the slowest GPU determines the synchronization barrier, producing a straggler effect and forcing faster GPUs to idle [2509.04719].

The defining idea of STADI is that diffusion inference exposes two schedulable axes. The first is temporal: different GPUs can be assigned different denoising step counts. The second is spatial: different GPUs can be assigned different patch sizes. STADI therefore adapts both steps and patches so that faster GPUs perform more work and slower or busier GPUs perform less. This differs from prior patch-parallel schemes that primarily split only the spatial dimension and assume near-uniform devices, equal patch sizes, equal step counts, and synchronized execution every step [2509.04719].

A common source of ambiguity is the phrase “spatio-temporal adaptive diffusion.” In STADI proper, “temporal” refers to denoising-step scheduling and “spatial” refers to patch allocation across GPUs. In DSTPP and STADMs, by contrast, “spatio-temporal” refers to the joint structure of events in time and space or to the joint spatial-temporal structure of EEG signals. The shared phrase signals a family resemblance in adaptive diffusion design, but not a shared problem setting [2305.12403] [2407.03089].

## 2. System model and optimization target

STADI formulates the deployment environment as a set of \(N\) GPUs, where each GPU \(i\) has a relative computational capability \(c_i > 0\) and a background utilization \(\rho_i \in [0,1]\). From these quantities and profiling, the framework derives an effective speed \(v_i\). The inference objective is to determine, for each GPU, a step count \(M_i\) and a patch size \(P_i\) so as to minimize total latency \(R_{\text{total}}\), with per-GPU latency written as \(R_i(M_i, P_i, c_i, \rho_i)\) [2509.04719].

This formulation isolates the central systems problem: if all GPUs are forced to execute the same number of steps on the same patch sizes, the slowest GPU determines the synchronization cadence. In heterogeneous systems this produces persistent workload imbalance. STADI interprets the problem not as a purely spatial partitioning task, but as a joint allocation problem in which temporal and spatial work must be balanced together [2509.04719].

The framework introduces a hybrid scheduler that combines temporal adaptation and spatial adaptation. A warmup phase is used before asymmetric step allocation is introduced. This warmup keeps the inference trajectory aligned across devices and is intended to avoid early divergence. After warmup, the scheduler departs from equal-step execution and assigns different step budgets according to effective speed [2509.04719].

A plausible implication is that STADI treats synchronization structure itself as part of the optimization problem. The target is not merely to reduce the arithmetic workload on slow GPUs, but to reduce idle periods, synchronization stalls, and execution irregularity induced by heterogeneous devices and background load.

## 3. Hybrid scheduling: step allocation and elastic patch parallelism

STADI’s temporal adaptation is centered on a computation-aware step allocator applied after warmup. All GPUs first execute the same warmup length \(M_{\text{warmup}}\). After that point, step counts are assigned according to effective speed relative to the fastest GPU speed \(v_{\max}\), using thresholds \(a\) and \(b\) with \(0 < b < a < 1\). Fast GPUs retain the full base step budget \(M_{\text{base}}\); moderately slower GPUs receive a reduced budget, approximately half the base steps blended with the warmup length; extremely slow GPUs may be excluded [2509.04719].

The paper characterizes this design as a least-common-multiple-minimizing quantization strategy. The key idea is that step counts should not vary arbitrarily, but should be quantized into simple ratios, especially halving. If one GPU executes \(M\) steps and another executes \(M/2\), their synchronization points remain aligned at structured intervals. The practical consequences described are fewer denoising steps on slower GPUs, fewer synchronization events, less communication overhead, less idle time for fast GPUs, and a reduced desynchronization burden [2509.04719].

After warmup, the inference behavior becomes asymmetric. Fast GPUs alternate between synchronizing with all-gather on some steps and reusing the previous buffer on other steps, while slower GPUs continue synchronized updates each step. This mechanism is intended to preserve sufficient alignment without forcing all devices into the same synchronization schedule throughout the full denoising trajectory [2509.04719].

Spatial adaptation complements temporal adaptation through elastic patch parallelism. STADI equalizes an effective computational rate across GPUs and sets patch sizes according to
$$
P_i = \frac{v_i/M_i}{\sum_{j=0}^{N-1} v_j/M_j} \cdot P_{\text{total}}.
$$
Each GPU therefore receives a patch size proportional to its effective processing rate \(v_i/M_i\): faster GPUs with larger effective throughput receive larger patches, while slower GPUs with reduced step counts receive smaller patches [2509.04719].

The framework presents this as a complementary spatial mechanism rather than an alternative to step adaptation. Temporal adaptation alone can still leave residual imbalance if patch sizes remain too large on slow devices. Elastic patch resizing is the repair mechanism that balances per-step compute time so that GPUs finish each synchronization interval at roughly the same time [2509.04719].

## 4. Algorithmic operation, theoretical rationale, and empirical performance

Algorithmically, STADI takes as input the initial noise \(x_{t_0}\), the diffusion model \(\epsilon_\theta\), the GPU speeds \(\{v_i\}\), the base step count \(M_{\text{base}}\), and the warmup step count \(M_{\text{warmup}}\). For each GPU, it computes \(M_i\) from the temporal scheduling rule and \(P_i\) from the elastic patch formula. Inference then proceeds with synchronous warmup, followed by an asymmetric phase in which faster GPUs may skip some synchronization or buffer updates while slower GPUs continue synchronized execution, after which the final denoised image is returned [2509.04719].

The paper supports the step-asymmetry design with a temporal redundancy bound for DDIM-style updates. For synchronized inference in the DistriFusion style, the difference between neighboring temporal states is bounded by \(\mathcal{O}(1/M)\). For STADI, when step counts differ but satisfy a structured relation such as \(nM_i = M_j = M\), the cross-device discrepancy is also bounded by \(\mathcal{O}(1/M)\). This is used as the theoretical rationale for the claim that asymmetric step schedules can preserve the same asymptotic alignment behavior when the schedule is quantized compatibly [2509.04719].

The reported experimental setup uses 2 RTX 4090 GPUs connected via PCIe, each with 24 GB VRAM, under Python 3.10, PyTorch 2.6.0, NCCL 2.21.5, and the diffusers library with custom distributed scheduling. The diffusion model is Stable Diffusion XL (SDXL), a 2.3B-parameter model, and the dataset is the COCO Captions 2014 validation set. Baselines are DistriFusion / patch parallelism and tensor parallelism, reimplemented with the same NCCL primitives for fairness. Evaluation uses end-to-end latency, speedup ratio, PSNR, LPIPS, and FID. Heterogeneity is simulated by running a GPU occupancy program before inference, with occupancy patterns \([0\%,20\%]\), \([0\%,40\%]\), \([0\%,60\%]\), and, under fixed total occupancy, \([35\%,45\%]\), \([30\%,50\%]\), and \([25\%,55\%]\). Default settings are \(M_{\text{base}} = 100\), \(M_{\text{warmup}} = 4\), \(P_{\text{total}} = 32\), \(a = 0.75\), and \(b = 0.25\) [2509.04719].

The main quantitative claim is that STADI reduces end-to-end latency by up to 45% compared with patch parallelism. The paper also reports latency reductions of 12%–45% in one occupancy family, 4%–39% when total occupancy is fixed but redistributed, and 27.01%–45.36% overall in heterogeneous settings. Tensor parallelism is consistently worst. Ablations separate SA, meaning spatial adaptation, from TA, meaning temporal adaptation. Baseline latencies without adaptation are 14.91 s, 19.41 s, and 23.67 s for \([0\%,20\%]\), \([0\%,40\%]\), and \([0\%,60\%]\), respectively. SA only yields speedups of 1.12×–1.34×; TA only yields 1.32×–1.82×; TA + SA achieves the best overall results at 1.37×, 1.70×, and 1.83× depending on occupancy. Quality is reported to degrade slightly when steps are reduced, but semantic quality remains strong, with FID differences versus the original reportedly less than 1 in many comparisons [2509.04719].

## 5. Relation to earlier spatio-temporal adaptive diffusion methods

Before the exact STADI acronym was used for heterogeneous multi-GPU scheduling, closely related spatio-temporal adaptive diffusion ideas appeared in model-centric settings. In "Spatio-temporal Diffusion Point Processes" [2305.12403], the target problem is a spatio-temporal point process in which each event is \(x_i = (t_i, s_i)\), the history is \(H_t = \{x_i \mid t_i < t,\ x_i \in S\}\), and the modeling objective is the conditional density \(p(t,s \mid H_t)\). The paper argues that time and space are strongly entangled, criticizes prior methods for assuming conditional independence or unilateral dependence, and proposes a denoising diffusion probabilistic model that decomposes the joint distribution into multiple small denoising steps. Its adaptive element is a spatio-temporal co-attention denoising network whose attention weights depend on both event history and diffusion step, allowing time-space dependence to vary across datasets and denoising stages. The method is called DSTPP rather than STADI, but it is explicitly identified as a close conceptual match to a spatio-temporal adaptive diffusion inference framework [2305.12403].

DSTPP’s empirical scope is broad: Earthquake in Japan, COVID-19 spread in New Jersey, Citibike urban mobility in NYC, Crime in Atlanta neighborhoods, HawkesGMM synthetic data, and a Synthetic-Independent dataset. The paper states that the framework outperforms state-of-the-art baselines remarkably, with an average improvement of over 50%, and also reports over 20% average reduction in spatial NLL compared with the second-best model. Its central interpretive claim is that co-attention weights adapt to whether the data are strongly entangled or nearly independent, as illustrated by strong mutual dependence on Earthquake and mostly self-focused attention on Synthetic-Independent [2305.12403].

A second near match appears in "Generative AI Enables EEG Super-Resolution via Spatio-Temporal Adaptive Diffusion Learning" [2407.03089]. There the method is named STADMs rather than STADI and addresses EEG super-resolution from low-resolution EEG to 256-channel high-resolution EEG. The problem is posed as learning a mapping from \(x \in \mathbb{R}^{C_l \times N}\) to \(y \in \mathbb{R}^{C_h \times N}\), with synthetic SR EEG \(y_{sr} = \mathit{F}(x,\theta)\). The architecture is a conditional latent diffusion model operating on latent representations produced by a pre-trained masked autoencoder. Its adaptive branch is the spatio-temporal condition module (STC), which extracts temporal correlations across time points and spatial correlations across EEG channels using known electrode positions. The reverse-process generator is the multi-scale Transformer denoising module (MTD), which uses multi-scale 1D convolution blocks with kernel sizes 3, 5, 7, and 9, followed by diffusion Transformer blocks with self-attention and cross-attention [2407.03089].

STADMs is evaluated on the Localize-MI dataset, comprising seven drug-resistant epilepsy patients and 61 sessions, recorded with the EGI NA-400 amplifier at 8000 Hz and preprocessed with MNE using high-pass filtering, notch filtering at 50, 100, 150, and 200 Hz, and 350 ms epoching. The low-resolution configurations are 128, 64, 32, and 16 channels derived from EGI-128, EGI-64, EGI-32, and EGI-16 montages. The paper reports that STADMs outperform Deep-CNN, WGLAE, DAE, and EEGSR-GAN on PCC, MAE, NMSE, and SNR; that synthetic SR EEG improves EEGNet-based classification across Accuracy, Precision, Recall, and F1; and that, under eLORETA, SR EEG reduces localization error and approaches HR EEG performance. As in DSTPP, the exact acronym differs, but the underlying pattern is a spatio-temporal conditioning mechanism that adaptively guides diffusion inference [2407.03089].

## 6. Assumptions, limitations, and interpretive boundaries

STADI’s analysis assumes bounded diffusion states, bounded model outputs, smooth diffusion parameters, effective GPU speeds that can be profiled offline, and current occupancy that can be measured via system APIs. It also assumes that the system can support the required communication primitives, including customized all-gather for uneven patch sizes. These assumptions matter because the framework’s scheduling decisions depend on accurate profiling and on communication support for nonuniform spatial partitioning [2509.04719].

Several limitations are stated explicitly. Extremely slow GPUs may need to be excluded. Uneven patch sizes require custom asynchronous communication. The current evaluation is on a 2-GPU setup rather than a large cluster. Quality can degrade slightly when step reduction is applied. Patch-size and step relationships may become non-linear under strong fixed overheads. Some operators constrain patch dimensions, for example through power-of-two-like requirements [2509.04719].

These boundaries distinguish STADI from an exact quality-preserving equivalence to synchronized homogeneous execution. The framework is a latency-quality tradeoff mechanism whose principal goal is to improve resource utilization and reduce bottlenecks in realistic heterogeneous deployments. The reported results indicate that the quality cost is small in the tested regime, but not zero [2509.04719].

Future directions identified in the paper are larger GPU clusters, synergy with model compression, and extension to video generation and other tasks. A plausible implication is that STADI’s core abstraction—joint scheduling over denoising steps and spatial partitions—could be generalized beyond the particular SDXL-on-two-GPU configuration used in the current evaluation, provided that the same assumptions about profiling, synchronization structure, and communication support remain tenable [2509.04719].

Source: https://www.emergentmind.com/topics/spatio-temporal-adaptive-diffusion-inference-stadi