Technical Summary: Align-RAG

Problem Statement

Retrieval-augmented forecasting (RAG) for Time Series Foundation Models (TSFMs) aims to adapt frozen models to new domains without fine-tuning by retrieving relevant past–future windows from a reference corpus. However, recent state-of-the-art methods (e.g., TS-RAG, TimeRAF, RATD) rely on the assumption that frozen backbones lack inherent mechanisms to dynamically incorporate retrieved context. Consequently, these approaches introduce learned fusion modules (e.g., adaptive mixers, prompting modules, or projection heads) trained to merge retrieved examples with the backbone's representation. This paper challenges the necessity of such learned components, positing that the performance gains attributed to them may instead stem from a lack of proper alignment between the retrieved demonstrations and the query.

Methodology: Align-RAG

The authors introduce Align-RAG, a training-free method that enhances a frozen TSFM by applying closed-form transformations to retrieved past–future windows before they enter the model's context. The method operates without modifying the backbone or training any new parameters.

Core Mechanisms

1. Amplitude Alignment: For each retrieved past window $p_i$, a regularized affine map is fitted to match the query past $q$. The slope $a_i$ and intercept $b_i$ are derived using a Wiener-style shrinkage estimator to prevent noise amplification on low-variance neighbors:
   $$a_i = \frac{\sigma_q \sigma_{p_i}}{\sigma_{p_i}^2 + (\sigma_q/M)^2}, \quad b_i = \mu_q - a_i \mu_{p_i}$$
   This map is applied to both the retrieved past and its corresponding future ($\tilde{p}_i, \tilde{f}_i$) to preserve the pair relationship.
2. Phase Alignment: An integer-lag shift $\tau^*_i$ is computed via cross-correlation to align the phase of the rescaled retrieved past with the query. This shift is applied to the future component as well, ensuring temporal consistency.
3. Context Construction:
   • Diversification: Maximal Marginal Relevance (MMR) selects a diverse subset of $K=10$ neighbors from a pool of 20 to avoid near-duplicates.
   • Token Budget: A tiered layout packs demonstrations into the model's context window (e.g., 2048 tokens for Chronos-Bolt), balancing detailed medium-range contexts with dense recent steps.
4. Prediction Fusion: The final forecast combines the model's in-context output with a distance-weighted prior of the aligned futures. A second forward pass on the query alone (without demonstrations) is performed, and the two predictions are averaged to stabilize the output.

Key Contributions

1. Training-Free Alignment: The introduction of Align-RAG, which demonstrates that aligning retrieved demonstrations in amplitude and phase is sufficient to unlock dynamic in-context learning in frozen TSFMs, eliminating the need for learned fusion adapters.
2. Empirical Superiority: Comprehensive evaluation showing Align-RAG outperforms the trained TS-RAG mixer on a frozen Chronos-Bolt backbone across all seven datasets of the standard benchmark (ETTh1, ETTh2, ETTm1, ETTm2, Weather, Exchange, Electricity), achieving an average MSE reduction of 3.75%.
3. Cross-Architecture Generalization: The method improves zero-shot performance on four additional frozen TSFMs (Chronos-2, TimesFM-2.0, Moirai, Toto) with gains ranging from 2.5% to 13.7% without per-backbone tuning.
4. Mechanistic Insight: Through ablation and behavioral analysis, the authors isolate that alignment (not retrieval ranking) drives the performance gains. They further demonstrate that aligned demonstrations induce prediction shifts in the frozen backbone that track a closed-form ridge predictor, suggesting the model performs implicit regression on the aligned pairs.

Results

• Benchmark Performance: On the Chronos-Bolt backbone, Align-RAG achieves lower MSE on 7/7 datasets and lower MAE on 6/7 datasets compared to TS-RAG. The improvement is statistically significant on small datasets via paired moving-block bootstrap.
• Ablation Studies:
  • Retriever vs. Alignment: Replacing the trained TS-RAG retriever with random retrieval while keeping alignment active still outperforms the trained TS-RAG mixer on 4/7 datasets. Conversely, removing alignment (even with a trained retriever) causes MSE to degrade by over 20 percentage points relative to zero-shot.
  • Component Analysis: Amplitude and phase alignment alone close most of the gap to the trained baseline; diversification and layout contribute further marginal gains.
• Mechanism Validation:
  • Implicit Ridge Regression: The cosine similarity ($\rho_{GD}$) between the model's prediction shift and a closed-form ridge predictor increases from 0.30 (unaligned) to 0.45 (aligned).
  • Future-Shuffle Control: Shuffling the futures of aligned pairs (breaking the $p \to f$ correspondence) reverses the MSE gains, confirming the model relies on the specific $(past, future)$ relationship rather than simply averaging retrieved futures.

Significance and Claims

The paper claims that the assumption requiring learned fusion modules for retrieval-augmented forecasting is unnecessary. The primary contribution is the demonstration that frozen TSFMs already support dynamic in-context use of retrievals, provided the demonstrations are properly aligned in amplitude and phase.

The authors argue that previous gains attributed to learned adapters were largely compensating for the absence of such alignment. Consequently, they propose that closed-form alignment should be the default baseline for retrieval-augmented forecasting before any fusion module is trained. The work suggests that the "alignment" of context is the critical factor, not the "learning" of a fusion mechanism, offering a simpler, parameter-free alternative that matches or exceeds state-of-the-art trained methods.