[ICML 2026] PATRA: Pattern-Aware Alignment and Balanced Reasoning for Time Series Question Answering

Table of Contents

• Installation
• Abstraction
• Framework
• Training Workflow
• Results
• Training Step
• Evaluation
• Demo

This code is a PyTorch implementation of our ICML'26 paper "PATRA: Pattern-Aware Alignment and Balanced Reasoning for Time Series Question Answering". [arXiv]

Installation

pip install -r requirements.txt

Abstraction

Time series reasoning demands both the perception of complex dynamics and logical depth. However, existing LLM-based approaches exhibit two limitations: they often treat time series merely as text or images, failing to capture the patterns like trends and seasonalities needed to answer specific questions; and when trained on a mix of simple and complex tasks, simpler objectives often dominate the learning process, hindering the development of deep reasoning capabilities. To address these limitations, we propose the Pattern-Aware Alignment and Balanced Reasoning model (PATRA), introducing a pattern-aware mechanism that extracts trend and seasonality patterns from time series to achieve deep alignment. Furthermore, we design a task-aware balanced reward to harmonize learning across tasks of varying difficulty, incentivizing the generation of coherent Chains of Thought. Extensive experiments show that PATRA outperforms strong baselines across diverse Time Series Question Answering (TSQA) tasks, demonstrating superior cross-modal understanding and reasoning capability.

Framework

The framework of PATRA, which contains a Text Encoder for obtaining text embedding from pre-trained representations, a TS Encoder for embedding the time series, a Pattern-Aware Alignment to align the gap between text embedding and time series embedding, and an LLM Backbone to reason the questions.

Training Workflow

Overview of PATRA’s two training stages. Alignment Stage performs supervised fine-tuning (SFT) with cross-entropy (CE) loss, while Reasoning-Enhanced Stage applies reinforcement learning, including format and task reward under the GRPO strategy to optimize the policy model.

Results

Training

We use Llama-factory to train PATRA. Scripts for the two stages are provided in train/scripts:

• Alignment Stage (SFT): bash train/scripts/train_patra-sft.sh
• Reasoning-Enhanced Stage (GRPO): bash train/scripts/train_patra_grpo.sh

Note: The PATRA-init checkpoint will be released coming soon. The training data (SFT & GRPO) is available at DecisionIntelligence/PATRA-TRAIN, and the final PATRA checkpoint at DecisionIntelligence/PATRA-7B.

Both shell scripts contain placeholders that must be replaced before launching:

Script	Placeholder	Replace with
train_patra_sft.sh	--model_name_or_path "base_model_path"	Path to the PATRA-init checkpoint (coming soon)
train_patra_sft.sh	--dataset "sft-all"	Registered SFT dataset name (points to sft.jsonl from PATRA-TRAIN via train/data/dataset_info.json)
train_patra_sft.sh	--output_dir "output_dir"	Directory to save the SFT checkpoint
train_patra_grpo.sh	--model_name_or_path "sft-model-path"	Path to the SFT checkpoint produced by the previous stage
train_patra_grpo.sh	--dataset "/path/to/grpo_dataset.jsonl"	Absolute path to grpo.jsonl from PATRA-TRAIN (loaded directly, not via dataset_info.json)
train_patra_grpo.sh	--output_dir "output_dir"	Directory to save the GRPO checkpoint

The SFT stage resolves --dataset through train/data/dataset_info.json, while the GRPO stage loads the file directly via datasets.load_dataset('json', ...). Pass a registered name for SFT and a .jsonl path for GRPO.

Evaluation

Evaluation splits are hosted on Hugging Face: DecisionIntelligence/PATRA-EVAL. Four TSQA subsets are provided — comprehension, recognition, reasoning, prescience — each row containing input (prompt with a <ts><ts/> placeholder), timeseries, question_format (multiple_choice / true/false / open_ended_question), and answer.

Pull them locally and run one task at a time:

huggingface-cli download DecisionIntelligence/PATRA-EVAL --repo-type dataset --local-dir eval/data

python eval/patra_eval.py \
    --task reasoning \
    --model_path DecisionIntelligence/PATRA-7B \
    --data_dir eval/data \
    --device cuda:0

--out_dir (defaults to eval/results/) is optional. The runner reports accuracy on multiple-choice / true-false questions and ROUGE-L on open-ended ones, and writes <task>-result.json to out_dir.

Demo

python demo/patra_demo.py

Note: demo/patra_demo.py loads the final PATRA checkpoint from DecisionIntelligence/PATRA-7B. Set model_path in the script to that repo id (or a local snapshot).

Citing PATRA

If you find this resource helpful, please consider to cite our research:

@inproceedings{lu2026patrapatternawarealignmentbalanced,
      title={PATRA: Pattern-Aware Alignment and Balanced Reasoning for Time Series Question Answering}, 
      author={Junkai Lu and Peng Chen and Xingjian Wu and Yang Shu and Chenjuan Guo and Christian S. Jensen and Bin Yang},
      booktitle={ICML},
      year={2026}
}