NeurIPS 2026 (accepted main-conference paper) Time Series spatial-temporal forecasting heterogeneity-aware experts spatial-temporal decoupling spatial mixing adaptive gating

AdaST: Adaptive Coupling for Spatial-Temporal Forecasting¶

Conference: NeurIPS 2026 (accepted main-conference paper)
arXiv: 2609.36119
Code: https://github.com/LzyFischer/AdaST
Area: Time Series
Keywords: spatial-temporal forecasting, heterogeneity-aware experts, spatial-temporal decoupling, spatial mixing, adaptive gating

TL;DR¶

AdaST decomposes spatial-temporal inputs into temporal-specific, spatial-specific, and jointly coupled representations, models them separately, and adaptively recomposes them with correlation-modulated gates, achieving the best reported results on four short-term sensor forecasting benchmarks and reducing PurpleAir MAE from the strongest comparator's 0.511 to 0.489.

Background & Motivation¶

Traffic, air-quality, and energy sensor data all have a temporal axis and a variable or location axis, but having both axes does not imply that they should always be strongly coupled. DCRNN, STGCN, and GWNet typically combine spatial interactions and temporal evolution in a fixed architecture: predictions at each location always receive some form of cross-location information. This suits congestion propagation, but can inject irrelevant information into series primarily governed by their own history. Conversely, independent temporal modeling misses short-term cross-location propagation caused by abrupt events. The paper distinguishes temporal-dominated, spatial-dominated, and strongly coupled regimes rather than treating every dataset as the same spatial-temporal process.

The difficulty is that these regimes are neither predefined dataset labels nor necessarily uniform within a dataset. Different traffic sensors may follow different patterns, and the same location may rely on different information during regular commuting and abrupt changes. The paper first generates synthetic data with three known mechanisms and compares temporal-only, spatial-only, and joint architectures, finding better performance when architecture and generation mechanism align. This is controlled evidence for the motivation, not proof that real-world data admits an identifiable causal decomposition. Spatial modeling introduces another trade-off: predefined graphs can restrict interactions, whereas per-sample spatial attention repeatedly scores all node pairs.

AdaST does not select one fixed architecture for an entire dataset. Instead, it retains three candidate information paths within each input, lets location and periodic context guide decomposition, and lets predictive representations determine each path's contribution. Here, decoupling means learning three representations with different architectural roles, not recovering unique physical signal components from observations. Core idea: use four heterogeneity-aware experts to provide multiple views of three correlation patterns, process them through role-aligned branches, and recompose them across time and locations with correlation-modulated gates rather than unconditionally strengthening unreliable cross-axis dependencies.

Method¶

Overall Architecture¶

The model takes observations from the past \(T\) steps at \(N\) locations, with \(D\) features per location, and predicts the corresponding values over the next \(T'\) steps. Observations are first projected into a hidden space. Four experts add location, time-of-day, day-of-week, and joint spatial-temporal context, respectively; every expert produces all three component types, rather than serving only one branch. Concatenating outputs of the same component type across experts creates three equally wide branch inputs.

The temporal component passes only through temporal attention, the spatial component only through spatial mixers, and the coupled component through alternating temporal and spatial operations. The recomposer computes sigmoid gate scores from the processed representations and adjacent-index cosine correlation measures within each component. It multiplies these quantities and then applies softmax across the three branches. The weighted representation passes through a temporal projection and an output layer. Learned gate scores therefore differ from the final mixture weights, and correlation measures are not merely post-prediction evaluation statistics.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    X["Historical observations<br/>and time indices"] --> E["Heterogeneous multi-view decomposition<br/>four experts, three components each"]
    E --> M["Role-aligned spatial-temporal modeling<br/>temporal / spatial / joint branches"]
    M --> R["Correlation-modulated recomposition<br/>sigmoid → modulation → softmax"]
    R --> P["Temporal projection<br/>and output layer"]
    P --> Y["Future multi-step predictions"]
    Y -.->|training only: prediction supervision| L["Training objective"]
    GT["Future ground-truth observations"] -.->|training only| L

Solid arrows show the forward prediction path, while dashed arrows indicate prediction supervision during training; future observations are not required at test time. The paper does not explicitly specify the exact training loss, so the diagram does not invent a loss formula.

Key Designs¶

1. Heterogeneous multi-view decomposition: every expert represents all three correlation patterns

A fixed decomposer can impose the same splitting rule on every location and time period. AdaST first linearly projects the observations into a base representation and separately concatenates four kinds of learnable context. The spatial expert uses node-specific embeddings broadcast across time. The time-of-day expert looks up embeddings using each observation's time-of-day index, while the day-of-week expert uses its weekday index; both periodic embeddings are expanded across locations. The joint spatial-temporal expert uses learnable embeddings with both temporal and location axes. Experts differ in the context they receive, not in whether they exclusively produce spatial, temporal, or joint outputs. Even the spatial expert produces a temporal component, because location identity can inform which historical patterns should be retained.

Each expert's projection head produces a representation of total width \(3D_H/4\), split into joint, temporal, and spatial components of width \(D_H/4\) each. Components of the same type are then concatenated along the feature axis: the temporal branch collects temporal outputs from all four experts, and the other branches follow the same rule. The following relation captures this concatenation, with \(k\in\{n,h,w,a\}\) denoting location, time-of-day, day-of-week, and joint experts, and \(c\in\{t,s,st\}\) denoting component types.

\[ [Z_k^{(st)};Z_k^{(t)};Z_k^{(s)}]=f_k(H_k),\qquad H^{(c)}=[Z_n^{(c)}\|Z_h^{(c)}\|Z_w^{(c)}\|Z_a^{(c)}]. \]

Thus, four experts contribute twelve intermediate blocks that form three inputs of width \(D_H\); the model neither selects one expert nor averages the experts' complete outputs. With the default \(D_H=256\), each block has width 64 and each expert outputs 192 dimensions in total. No explicit reconstruction constraint requires the original observations to equal the sum of the three components. These are predictive hidden representations, not a strict additive decomposition of observable signals.

2. Role-aligned spatial-temporal modeling: accessible information defines each component's role

Naming components does not ensure that they learn different information. AdaST further constrains information exchange through branch structure. The temporal branch applies multi-head self-attention along the time axis within each location, with residual connections, dropout, layer normalization, and a feed-forward network; this branch does not mix nodes. The spatial branch exchanges information along the node axis through a spatial mixer without invoking temporal attention. The joint branch alternates temporal attention and spatial mixing within each layer. All three branches stack \(L\) layers and retain equal output shapes for subsequent mixing across time and locations. “Spatial-specific” means that this branch does not perform dedicated temporal attention; it does not mean that the complete model discards historical inputs or temporal context.

The spatial mixer learns an \(N\times N\) assignment matrix, applies softmax to mix node representations, and then uses a residual connection and feed-forward network. In the paper's transpose notation, the transpose places the node axis in the position required for multiplication. Here, \(A\) is a learnable parameter, not a mandatory predefined road adjacency matrix.

\[ \operatorname{SpatMix}(H)=\operatorname{FFN}\!\left(H+(\operatorname{softmax}(A)H^\top)^\top\right). \]

The mixer still supports global node interactions, but its mixing matrix is not recomputed from query/key representations for every sample, avoiding per-sample spatial attention scoring and attention-map storage. This saving should not be confused with a lower asymptotic order: Appendix B explicitly gives both operators the same \(\mathcal{O}(N^2TD)\) aggregation cost, rather than linear complexity. Adaptivity primarily comes from input representations and subsequent mixture weights, so this matrix should not be described as a per-sample dynamic graph. Input-independent interactions may reduce noise, but removing spurious correlations remains an experimental interpretation rather than a consequence of the module's name.

3. Correlation-modulated recomposition: assess within-component consistency before normalizing branch contributions

After branch modeling, each branch uses a linear projection and sigmoid to produce gate scores indexed by historical time and node. Learned gates alone may still overemphasize unreliable representations, so the authors add cosine-based correlation measures. Temporal correlation averages cosine similarity between adjacent slices along time; spatial correlation averages similarity between adjacent-index slices along the node axis. The joint component averages both measures computed on its own representation. The equations below omit shared indices in the prose notation and express the adjacent-slice averaging mechanism of original Equations (16)–(18).

\[ \mathcal{C}_t(H)=\frac{1}{T-1}\sum_{i=1}^{T-1}\operatorname{CosineSim}(\hat H_{:,i,:,:},\hat H_{:,i+1,:,:}),\qquad \mathcal{C}_s(H)=\frac{1}{N-1}\sum_{j=1}^{N-1}\operatorname{CosineSim}(\hat H_{:,:,j,:},\hat H_{:,:,j+1,:}), \]

\[ \mathcal{C}_{st}(H^{(st)})=\frac{\mathcal{C}_s(H^{(st)})+\mathcal{C}_t(H^{(st)})}{2}. \]

“Spatial adjacency” is particularly easy to misread: Equation (17) uses node indices \(j\) and \(j+1\), without selecting neighbors by graph edges, geographical distance, or road connectivity. This differs from Appendix A's synthetic-data analysis, which computes Pearson correlation over spatially connected node pairs. Correlation is measured on learned component representations, not directly on raw observations as a physical coupling strength. The original equations also introduce a batch axis without fully specifying cosine reduction and broadcasting details, so this note does not infer all implementation dimensions.

The correlation measure raised to \(\alpha\) multiplies the sigmoid score. Softmax across the three components then produces nonnegative mixture weights summing to 1. The three equally shaped representations are combined with weights broadcast along features, and only then passed to the temporal projection and output mapping.

\[ g^{(c)}=\sigma(W_c\hat H^{(c)}),\qquad \tilde g^{(c)}=g^{(c)}[\mathcal{C}_c(H^{(c)})]^\alpha,\qquad w^{(c)}=\frac{\exp(\tilde g^{(c)})}{\sum_{r\in\{t,s,st\}}\exp(\tilde g^{(r)})}. \]

The three correlations are therefore not directly treated as mixture probabilities, nor are the modulated scores simply divided by their sum. The default \(\alpha=0.1\) controls the influence of correlation on gate scores. The Section 3 overview calls this “Correlation-Regularized Recoupling,” whereas Section 3.3 calls it “Correlation-Informed Recomposition.” The actual equations specify forward gate modulation, not a separate supervised correlation loss. Cosine similarity can be negative, but the paper does not explain how a non-integer power handles negative values or whether clipping or shifting is used. Appendix E's reference to “normalized correlation measures” does not define a transformation sufficient to resolve this issue; absolute values or truncation must not be assumed.

A Worked Example¶

Consider the default traffic forecasting setting: 12 historical observations at 5-minute intervals predict the next 12 steps, mapping one hour of history to one hour of forecasts. This example illustrates tensor flow without inventing unreported gate values.

At each historical step and node, the base observation representation reaches all four experts. Each adds its own 24-dimensional context and produces three 64-dimensional blocks. Four temporal blocks form a 256-dimensional temporal input, and the spatial and joint inputs also have 256 dimensions each. A sensor therefore retains time-of-day, location identity, and joint context views simultaneously; receiving information from the spatial expert does not remove its temporal component.

Each branch undergoes 3 layers of role-aligned modeling. The temporal branch reads the sensor's history, the spatial branch exchanges information through the learned node mixing matrix, and the joint branch performs both operations. The recomposer creates three-way weights on historical representations. If the learned representation indicates that stable history is more reliable, the temporal contribution can increase relatively, without forcing the spatial contribution to zero. The resulting weighted representation maps to the next 12 steps; gating does not directly assign an independent road graph to every future prediction step.

Loss & Training¶

The paper specifies Adam, an initial learning rate of 0.001, exponential decay, and batch size 16, with an NVIDIA A100 80GB GPU. All four expert embeddings have width 24, the hidden width is 256, there are 3 layers, and temporal attention uses 4 heads. Input and output lengths are both 12. Appendix B searches \(\alpha\) over \(\{0.01,0.05,0.1,0.5,1.0\}\) on the validation set and selects 0.1.

Data use Z-score normalization and a 60%/20%/20% training/validation/test split. PEMS is sampled every 5 minutes. PurpleAir's original 2-minute observations are resampled to 6 minutes, so its 12-step horizon corresponds to 72 minutes rather than one hour. Time-of-day tables cover 288 and 240 slots, respectively, and the weekday table covers 7 days.

Neither the main text nor the appendices explicitly states the mathematical optimization objective. Reporting MAE as a metric therefore does not justify claiming that the authors use an MAE training loss, and orthogonality, reconstruction, or correlation regularization losses must not be invented. What is clear is that decomposition, branch modeling, and recomposition serve prediction; correlation modulation is a forward mechanism that requires no additional coupling labels.

Key Experimental Results¶

Main Results¶

Table 2 compares four benchmarks against 16 methods. Lower is better for all metrics, averaged over the next 12 prediction steps. The strongest MAE comparator is selected separately for each dataset and is not necessarily strongest on every metric.

Dataset	AdaST MAE	Strongest comparator MAE	AdaST RMSE	AdaST MAPE
PurpleAir	0.489	NBeats: 0.511	0.994	22.43%
PEMS04	18.28	D2STGNN / HimNet: 18.32	29.87	12.17%
PEMS07	19.16	STAEformer: 19.46	32.54	8.02%
PEMS08	13.45	HimNet: 13.52	23.17	8.85%

PurpleAir's relative MAE improvement is approximately 4.3%, consistent with the main text. PEMS04 improves by only 0.04, so gains should not be characterized as large on every dataset. Some cached values concatenate rendered text and LaTeX duplicates; this table recovers identical values only where column order makes them unambiguous, without adding precision.

Appendix Table 5 establishes an important boundary: AdaST achieves MAE 0.407 on ETTh1 and 3.128 on METR-LA, but scores 0.083 and ranks 4th on ExchangeRate, behind PatchTST's 0.073. Its average rank over five datasets is 1.6. Thus, “best on all four main benchmarks” does not mean “best on every dataset,” and these additional experiments still use 12-step short-term forecasting.

Ablation Study¶

The following selection from Table 3 reports MAE and PEMS07 MAPE with the full model as a shared reference. The text defines “w/o gates” as uniform averaging; “w/o correlation” retains learned gates but removes correlation modulation.

Config	PurpleAir MAE	PEMS07 MAE	PEMS07 MAPE
full model	0.489	19.16	8.02%
Remove location expert	0.519	19.44	8.49%
Remove time-of-day expert	0.495	20.39	8.68%
Remove day-of-week expert	0.498	19.21	8.07%
Remove joint spatial-temporal expert	0.516	20.16	11.26%
Replace mixer with spatial attention	0.518	19.35	8.11%
w/o correlation	0.513	19.85	10.58%
w/o gates	0.526	20.29	9.57%

By MAE, the location expert is the most important of the four experts on PurpleAir, and the time-of-day expert is most important on PEMS07. By PEMS07 MAPE, however, removing the joint expert is worse at 11.26%. The paper's “largest degradation” statements must therefore be interpreted with a specified metric rather than as universal rankings.

Spatial mixer efficiency comes from Appendix Table 4, in seconds per epoch. Speedups retain the authors' rounded values and do not indicate a change in asymptotic complexity.

Dataset	Spatial mixer	Spatial attention	Reported speedup
PEMS04	233	341	1.5×
PEMS07	1154	2160	1.9×
PEMS08	141	220	1.6×
PurpleAir	66	134	2.0×

Key Findings¶

Table 3 supports contributions from both gating and correlation modulation. However, in Appendix Table 6, the spatial-attention replacement and full model both achieve PEMS04 MAE 18.28; not every replacement strictly worsens every metric.
Figures 3–4 show mixture-weight differences across datasets, time periods, and locations. These are internal model diagnostics, not validation against independent ground-truth coupling labels.
Figure 5's t-SNE separation supports differences among branch representations but does not prove causal disentanglement. The checklist explicitly states that multi-seed error bars are absent, leaving small margins without statistical significance evidence.

Highlights & Insights¶

Four experts and three role-aligned branches form a crossed design: heterogeneous context determines how inputs are viewed, while branch architecture determines how information can be exchanged. Separating these choices is more flexible than assigning each context type to one branch.
Recomposition uses both learned predictive gate scores and within-representation consistency. The transferable idea is reliability-modulated fusion, not treating adjacent-index cosine similarity as a universal measure of spatial structure.

Limitations & Future Work¶

The authors acknowledge additional parameters and potential memory costs from three branches, and evaluation is limited to fixed 12-step inputs and outputs. Large node sets and long-term forecasting remain unvalidated.
Spatial mixing retains quadratic node cost. Location embeddings and a fixed-size mixing matrix also do not automatically guarantee transfer to unseen nodes; further design and experiments are needed.
Node order affects the spatial correlation definition, while negative cosine powers and the training loss remain underspecified. Node-permutation tests, explicit implementation conventions, and multi-seed evaluation should precede stronger interpretability claims.

vs D2STGNN: D2STGNN separates temporal and spatial-temporal components. AdaST additionally retains a spatial-specific component and uses correlation-modulated gating for recomposition; the distinction concerns component roles and fusion rather than merely adding a prediction head.
vs GWNet / STAEformer: AdaST can reduce a path's relative contribution instead of only strengthening spatial or temporal modeling. Its input-independent spatial mixing weights save computation but sacrifice the flexibility of per-sample spatial scoring.
vs PatchTST / DLinear: ExchangeRate results show that a complex adaptive model can still lose to temporal methods. Task suitability and simple baselines cannot be replaced by the intuition that three branches are more comprehensive.

Rating¶

Novelty: 4/5 — A clear combination of three role-aligned components, heterogeneous views, and correlation modulation, built from familiar modules.
Experimental Thoroughness: 3/5 — Multiple benchmarks, ablations, and efficiency measurements, but no multi-seed or long-term forecasting validation.
Writing Quality: 3/5 — Clear motivation, with insufficient specification of recomposition terminology, negative correlations, and the training objective.
Value: 4/5 — A reusable adaptive fusion approach for short-term multivariate forecasting, showing why strong coupling should not be imposed in every setting.