STReasoner: Empowering LLMs for Spatio-Temporal Reasoning in Time Series via Spatial-Aware Reinforcement Learning

AI summary

Overview

• Research area: Natural language processing and large language model reasoning applied to time series data, specifically spatio-temporal reasoning over graph-structured numerical signals.
• Technical level: Advanced. The paper assumes familiarity with reinforcement learning (GRPO/PPO-style policy optimization), stochastic differential equations, graph neural modeling, and multi-modal LLM training pipelines.
• Scope: The paper defines spatio-temporal reasoning in time series as a formal task, builds a synthetic data generation pipeline and benchmark for it, and trains an 8B open model (STReasoner) with a spatial-aware reinforcement learning objective that beats proprietary models at a fraction of the cost.

What This Paper Is About

Time series in real systems like traffic networks, power grids, and disease spread are coupled across space and time, but existing models only predict future numbers rather than explaining what happened, where, when, and why. Answering questions such as "which upstream node caused congestion at Node 2 at 9:00?" requires jointly reasoning over numerical signals, an explicit graph of spatial dependencies, and a natural language query. The paper's goal is to make LLMs perform this kind of explicit, spatially grounded reasoning, and to build the data, benchmark, and training method needed to measure and improve it.

Key Contributions

1. A formal problem definition and data synthesis pipeline. The authors define spatio-temporal reasoning in time series as a mapping from a query, a set of node-wise time series, and a graph to an intermediate reasoning trace and a final answer. They build a six-agent pipeline grounded in network stochastic differential equations (Network SDEs) that generates time series with controllable drift, diffusion, time-varying adjacency weights, and propagation delays, along with aligned natural language descriptions.

2. ST-Bench, a four-task benchmark. The benchmark decomposes spatio-temporal reasoning into etiological spatial reasoning (T1), spatial entity identification (T2), spatial correlation reasoning (T3), and in-context spatio-temporal forecasting (T4). A companion dataset, ST-Align, provides 153,700 basic temporal, spatial, and spatio-temporal QA pairs for alignment pretraining.

3. STReasoner, a unified spatio-temporal reasoning model. A time series encoder (patchify plus a five-layer MLP) produces embeddings that are interleaved with text tokens in node order and fed to an LLM, with value-preserving normalization to retain numerical precision. The model is trained in three stages: large-scale alignment, supervised fine-tuning on rejection-sampled chain-of-thought data, and reinforcement learning.

4. S-GRPO, a spatial-aware reinforcement learning objective. An extension of Group Relative Policy Optimization that assigns an extra reward bonus only when the model performs better with the graph structure present than with it removed, creating a contrastive signal that rewards genuinely spatial reasoning strategies rather than superficial temporal pattern matching.

Main Findings

• STReasoner-8B leads on all four benchmark tasks. It reaches 95.65% accuracy on etiological reasoning, 75.71% on entity identification, 87.12% on correlation reasoning, and 65.59 MAE on forecasting, beating GPT-5.2 and Claude-4.5-Sonnet on the first three tasks while using roughly 0.004× the estimated cost.

• Large gains at minimal cost. Against the best baselines, STReasoner improves accuracy by an average of 17% on T1, 135% on T2, and 40% on T3, while remaining within 2.6% on T4 forecasting. Token usage is far lower (0.08M–0.72M input tokens versus up to 10.26M for text-prompted baselines).

• Strong zero-shot generalization to real data. On a causal QA dataset built from CausalRivers, STReasoner achieves 98.82% accuracy without any fine-tuning on that data, versus 83.18% for Claude-4.5-Sonnet and 22.32% for GPT-5.2, at a cost of $0.05 compared with $2.35 and $6.15.

• A dedicated time series encoder matters more than prompting modality. Text-prompted and image-prompted variants trained with the same pipeline both underperform the encoder-based model. Image prompting helps on shape-driven tasks but hurts on forecasting, which requires precise numerical values.

• All three training stages contribute. Alignment alone performs poorly (3.38% on T1), SFT alone reaches 90.34%, and SFT without alignment and RL still lags; alignment adds roughly 6.3% on average when combined with SFT and S-GRPO. Skipping SFT and applying S-GRPO directly collapses to 47.34% on T1 due to sparse reward signals under weak initialization.

• S-GRPO beats vanilla GRPO. Replacing GRPO with S-GRPO yields a 5.10% average improvement across tasks, and a substantially higher rate of responses that explicitly invoke spatial information.

• The model learns to depend on the graph over the course of training. On Task 1, the accuracy gap between inputs with and without graph structure grows from 0.5 points at RL step 1 to 6.2 points at step 51, showing progressive reliance on spatial structure rather than just general accuracy gains.

• Method gains exceed scale gains. STReasoner-8B outperforms an untrained-style Qwen3-14B baseline on multiple tasks, and STReasoner improves over the Qwen3 baseline at every size from 4B to 14B.

• Spatial reward magnitude is sensitive. Intermediate values (α = 0.1 and 0.2) perform best; large values (α = 1, 0.5) let the spatial bonus dominate correctness, and α = 0 removes spatial grounding entirely.

Methodology in Plain English

The work proceeds in four connected steps. First, the authors generate training data rather than collecting it. Because no existing datasets pair graph-structured time series with natural language explanations, they simulate systems using network stochastic differential equations, where each node has its own drift and diffusion terms and receives influence from neighbors through a time-varying adjacency matrix with explicit propagation delays. Six language-model agents coordinate this: one invents a scenario in prose, one parses it into structured specifications, one instantiates the SDE parameters, one assigns dynamic edge weights and delays, and two judge agents validate the scenario logic and the simulated parameters respectively, feeding corrections back. Simulation then produces the numerical series. From these, templated questions and answers are generated for alignment and reasoning tasks, and all 1,200 synthesized samples were manually inspected, leaving 1,064.

Second, the model architecture keeps time series and text separate until they meet inside the LLM. Series are split into patches and encoded by a small MLP; those embeddings are inserted as placeholder tokens at the positions where each node's series belongs in the text prompt, alongside the graph description and the question. This lets the LLM attend over numerical and textual information in one sequence.

Third, training happens in three stages. A large alignment stage teaches the model to associate text descriptions with time series properties using 153,700 basic QA pairs. A supervised fine-tuning stage then teaches reasoning, using chain-of-thought traces that were sampled from Claude-4.5-Sonnet and kept only when they produced the correct answer. Finally, reinforcement learning refines behavior.

Fourth, the RL stage uses the novel S-GRPO reward. For each question, the model generates one group of responses with the graph structure included and another with it removed. A response earns an extra bonus only if its reward with the graph exceeds a threshold multiple of the reward without it, so the model is rewarded specifically for reasoning that depends on spatial structure. Advantages are then computed within each group and the policy is updated with a clipped objective and a KL penalty against a reference model, following the GRPO formulation.

Why This Matters

This paper shifts attention in time series AI from point prediction toward explicit, interpretable reasoning that can support decisions. It provides the first structured benchmark for decomposing spatio-temporal reasoning into distinct capabilities, and it shows that a carefully trained 8B open model can outperform far larger proprietary systems on this class of problems at a small fraction of the cost.

Real-world applications:

• Traffic management: tracing which upstream road segment caused a downstream congestion event, accounting for travel delay and attenuation along the network.
• Epidemic surveillance: identifying source regions of disease spread and explaining multi-hop transmission paths across connected populations.
• Power grid resilience: assessing cascading failure risk under extreme weather by reasoning over grid topology and time-lagged propagation.
• Flood forecasting: following flood propagation along river networks to determine which upstream gauge readings predict downstream inundation.

Industry relevance: The cost profile is the headline for deployment. STReasoner operates at roughly 0.004× the estimated cost of proprietary models while outperforming them on three of four tasks, which makes on-premise or edge deployment of spatio-temporal reasoning feasible for operators of infrastructure networks who cannot send sensitive grid, traffic, or public health data to external APIs. The S-GRPO technique is also domain-agnostic and could be transferred to any reinforcement learning setup where a structural input such as a graph, map, or knowledge base should be genuinely used rather than ignored.

Future Directions

• Broaden beyond synthetic data. The authors acknowledge that their evaluation rests on simulated data plus a single real-world zero-shot test. Expanding to diverse real-world spatio-temporal datasets with natural language annotations is the clearest next step.
• Improve the time series encoder. A five-layer MLP is sufficient for the structured signals studied here, but richer architectures and tighter multimodal fusion may be needed for noisy, irregular, or high-dimensional real systems.
• Tune and generalize the spatial reward. The sensitivity analysis shows performance depends non-trivially on the reward magnitude α, and the contrastive mechanism assumes a meaningful "without structure" ablation exists. Extending S-GRPO to other structured inputs and finding more robust reward shaping are open problems.
• Scale responsibly. The reported training consumed 215.77 GPU-hours and 30.45 kg CO2e; extending to larger backbones and datasets raises efficiency and sustainability questions that the authors flag explicitly.

Target Audience

Researchers in NLP and time series modeling who work on LLM reasoning, tool-augmented or multimodal language models, or graph-based forecasting will get the most from this paper. It is also relevant to reinforcement learning researchers interested in reward designs that enforce grounding in structured inputs, and to practitioners in transportation, energy, epidemiology, and environmental monitoring who need interpretable, low-cost reasoning over sensor networks rather than raw numeric forecasts. Readers without background in RL policy optimization or stochastic differential equations will find the methodology sections dense and should start with the introduction, task definitions, and results tables.

Authors’ abstract

Spatio-temporal reasoning in time series involves the explicit synthesis of temporal dynamics, spatial dependencies, and textual context. This capability is vital for high-stakes decision-making in systems such as traffic networks, power grids, and disease propagation. However, the field remains underdeveloped because most existing works prioritize predictive accuracy over reasoning. To address the gap, we introduce ST-Bench, a benchmark consisting of four core tasks, including etiological reasoning, entity identification, correlation reasoning, and in-context forecasting, developed via a network SDE-based multi-agent data synthesis pipeline. We then propose STReasoner, which empowers LLM to integrate time series, graph structure, and text for explicit reasoning. To promote spatially grounded logic, we introduce S-GRPO, a reinforcement learning algorithm that rewards performance gains specifically attributable to spatial information. Experiments show that STReasoner achieves average accuracy gains between 17% and 135% at only 0.004X the cost of proprietary models and generalizes robustly to real-world data.