TS-RAG:検索拡張で時系列基盤モデルのゼロショット予測を強くする

電力需要、気温、為替レート、工場の多変量センサー。こうした時系列を「これから24時間どうなるか」と予測したい場面は山ほどあります。最近は、巨大なデータで事前学習した時系列基盤モデル(TSFM: Time Series Foundation Model)——ChronosやTimesFM、Moiraiなど——が登場し、追加学習なしのゼロショットでかなり当たる予測を返すようになりました。

ところが、これらのモデルには弱点があります。事前学習で見たパターンの「平均像」をうまく出力する一方で、非定常な揺らぎや分布シフト——たとえば「いつもと違うレジームに切り替わった直後」——には弱いのです。モデルの重みは凍結されているので、目の前のデータに合わせて動的に知識を取り込む手段がありません。

ここで効いてくるのが、大規模言語モデルで大成功したRAG(Retrieval-Augmented Generation、検索拡張生成)の発想です。LLMが「関連する文書を検索してプロンプトに足す」ことで賢くなるなら、時系列でも「いま予測したい直近の窓に似た過去の区間を検索し、その続きを手がかりにする」ことができるはずです。これを実現したのが、本記事で解説する TS-RAG(Ning et al., “TS-RAG: Retrieval-Augmented Generation based Time Series Foundation Models are Stronger Zero-Shot Forecaster”, NeurIPS 2025, arXiv:2503.07649)です。

TS-RAGを理解すると、次のような場面で武器になります。

  • エネルギー需要・気象予測: 季節やイベントで急変する系列に、過去の似たパターンの続きを当てて予測を補正する
  • 金融・為替の短期予測: 分布シフトの大きい系列で、似た局面の履歴から外挿を補強する
  • 多変量センサーの予測保全: 凍結した基盤モデルを再学習せず、知識ベースの差し替えだけで現場に適応させる

本記事の内容

  • TS-RAGの直感 — なぜ「似た過去の続きを検索する」ことが予測を助けるのか
  • アーキテクチャ:リトリーバ(埋め込み + 近傍検索)、知識ベースの作り方、統合モジュール ARM の数式
  • 評価データセットと主要結果(ETT/Weather/Electricity/Exchangeで最大 +6.84%)のサーベイ
  • Pythonによる簡易デモで機構を実証(論文のSOTA再現は主張しません)

前提知識

この記事を読む前に、以下を読んでおくと理解が深まります。

画像なし
埋め込みと類似度
ベクトル埋め込みと類似度の基本。TS-RAGの『意味的に近い過去窓を探す』の土台です。
画像なし
ベクトルデータベースとFAISS
近傍検索の実装。TS-RAGの知識ベースもこの仕組みでtop-kを引きます。
画像なし
ハイブリッド検索
検索でどう関連情報を引くか。RAGの検索パートの理解に役立ちます。

TS-RAGとは — 「似た過去の続き」を予測のヒントにする

TS-RAG 概念模式図:クエリ窓を埋め込み検索し似た過去の続きで予測を補正

まずイメージから入りましょう。あなたが明日の電力需要を予測したいとします。手元には、直近の需要の波形(これをクエリ窓と呼びます)があります。素朴なやり方は「この波形を基盤モデルに入れて、続きを出力してもらう」ことです。これがベースのTSFM予測です。

TS-RAGはここにもう一手間かけます。「過去の膨大な履歴の中から、いまのクエリ窓と形がそっくりな区間を探し出す」のです。そして見つけた似た区間について、「その区間の”続き”が実際どうだったか」を取り出します。似た出だしの波形は似た続き方をしやすいので、これは予測の強力なヒントになります。

上の概念図がその流れです。クエリ窓を埋め込みエンコーダでベクトルに変換し、知識ベース(過去窓とその続きをペアで貯めた倉庫)から、埋め込みが近い上位 $k$ 件の過去窓を近傍検索します。引いてきた「続き」たちを、ARM(Adaptive Retrieval Mixer)という統合モジュールが重み付けして混ぜ、凍結した基盤モデルのベース予測に足し込みます。これで「いつもの平均像」だけでなく「いまの局面に似た実例の続き」が反映され、非定常・分布シフトに強いゼロショット予測になります。

ここで効いているのは、RAGの考え方そのものです。モデルの重みを一切いじらず、外部の知識(過去の実例)を検索して動的に持ち込むことで、凍結モデルを「いまのデータに適応」させているわけです。では、この検索と統合が実際にどんな部品でできているのか、順に分解していきましょう。

アーキテクチャの全体像 — 3つの部品

TS-RAGは、論文のFigure 2が示すとおり3つの部品から成ります。論文ではFigure 1が枠組み全体の俯瞰(凍結したTSFMに検索した知識を足す、というRAGの発想図)を、Figure 2がその具体的なデータフロー(検索した上位 $k$ ペアの未来区間と入力埋め込みをARMで統合する内部配線)を示しています。ここでは後者の詳細図を引きます。

TS-RAG モデルアーキテクチャ(論文Figure2)

出典: Ning et al., “TS-RAG”, NeurIPS 2025, Fig.2

  1. ベース時系列基盤モデル(TSFM、凍結): 入力時系列を埋め込み・特徴抽出し、最後に予測値へ射影する。論文ではこのTSFMを (a)エンコーディング層(Instance Norm・パッチ化・Embedder)、(b)バックボーン(GPT/T5/Llama等のTransformer)、(c)射影層(MLP)の3層構造として一般化します。実験では主に Chronos-Bolt をバックボーンに採用。Chronos-Boltは前身のChronosと違い、離散トークン化をやめてパッチ入力+デコーダ表現から複数ステップの分位点予測を一気に出す設計で、推論が速く精度も高いのが採用理由です。
  2. リトリーバ(Retriever): バックボーンとは別の事前学習済み時系列エンコーダ(論文ではChronosのエンコーダ)を流用し、クエリ窓と知識ベース内の過去窓を同じ埋め込み空間に写して、ユークリッド距離が近い上位 $k$ 件を引く。エンコーダ自体は学習しません。
  3. ARM(Adaptive Retrieval Mixer)augmentation module: 引いてきた「続き」たちと、バックボーンの内部表現を、多頭注意とゲートで動的に統合する。学習するのはこのARMと小さな射影器 $f_{\text{MLP}}$ だけで、バックボーンとリトリーバ・エンコーダは凍結したままです。

図で凍結(雪マーク)と学習可能(炎マーク)が色分けされている点に注目してください。学習対象を最小限に絞ることで、巨大な基盤モデルを再学習せずに済み、ゼロショット適応が現実的なコストで回ります。図2の流れを言葉で追うと、入力時系列 → リトリーバが知識ベースから上位 $k$ ペア(過去窓と対応する未来区間)を取得 → 取得した未来区間を射影器で埋め込み → 入力埋め込みと連結してARM(MHA → FFN → 重み付け → 混合)へ → TSFMの出力射影層で最終予測、となります。

この設計のうれしさは、知識ベースを差し替えるだけで適応先を変えられることです。同じ凍結モデルでも、知識ベースを電力データにすれば電力に、気象データにすれば気象に寄せられます。では、その知識ベースはどう作るのでしょうか。

知識ベースの作り方 — 「過去窓」と「その続き」をペアで貯める

TS-RAGの心臓部は知識ベースです。ここには、過去の時系列を一定長で切り出した(文脈窓, その続き)のペアを大量に貯めます。

論文では、各ペアを次のように定義します。

$$ \begin{equation} D = \{(x_i, e_i, y_i) \mid i = 1, 2, \dots, n\} \end{equation} $$

ここで $x_i \in \mathbb{R}^T$ は長さ $T$ の文脈窓(過去の区間)、$y_i \in \mathbb{R}^L$ はその続き(長さ $L$ の未来区間)、$e_i \in \mathbb{R}^d$ は事前学習済みエンコーダ $f_{\text{enc}}$ で $x_i$ を変換した埋め込みベクトルです。$n$ はペアの総数。つまり知識ベースは「過去窓・その埋め込み・その続き」の三つ組の集合です。

ポイントは、埋め込みをあらかじめ計算して一緒に貯めておくことです。こうすれば、推論時はクエリ窓を一度だけエンコードし、貯めてある埋め込みと距離を測るだけで検索が回ります。エンコーダを毎回叩き直す必要がありません。

論文では、知識ベースの中身として Chronos の事前学習データ(TSMixupで多ドメインの時系列を混ぜ合わせたもの)から一様サンプリングした 500万点を使い、これを文脈長で区切って 約280万ペアを構築しています。さらに ARM 自体の学習には別に 2600万ペアを用意します。多ドメインの実例を幅広く貯めることで、未知のデータにも「似た出だし」が見つかりやすくなる狙いです。

知識ベースができたら、いよいよ「クエリ窓に似た過去窓を引く」検索の番です。

リトリーバ — 埋め込みの近さで上位 k 件を引く

リトリーバの仕事は、クエリ窓 $x_q$ に意味的に近い過去窓を知識ベースから選ぶことです。生の波形どうしを直接比べると、スケールやノイズに振り回されます。そこで、いったん埋め込み空間に写してから距離を測ります。

まず、クエリ窓を事前学習エンコーダで埋め込みます。

$$ \begin{equation} e_q = f_{\text{enc}}(x_q) \end{equation} $$

次に、この $e_q$ と知識ベース内の各埋め込み $e_i$ とのユークリッド距離を計算します。

$$ \begin{equation} d(e_q, e_i) = \| e_q – e_i \|_2, \quad \forall i \in \{1, 2, \dots, n\} \end{equation} $$

そして距離が小さい順に上位 $k$ 件を選びます。

$$ \begin{equation} C = \mathrm{TopK_{min}}\big(\{(x_i, y_i, d(e_q, e_i)) \mid i = 1, \dots, n\},\ k\big) \end{equation} $$

$\mathrm{TopK_{min}}(\cdot)$ は距離 $d(e_q, e_i)$ の小さい順に $k$ 件を返す操作です。返ってきた集合 $C$ には、最も関連の高い「(過去窓, その続き)ペア」が $k$ 件入っています。実際に予測の補正に使うのは、このうちの続き $y_i$ の部分です。

なぜユークリッド距離で十分なのか、と思うかもしれません。鍵は埋め込みエンコーダ側にあります。エンコーダは正規化と特徴抽出を通じて、波形の形状(レジーム)を捉えた表現を作ります。だから埋め込み空間での近さが「波形の似ている度合い」に対応し、単純な距離でも意味のある検索になるのです。なお実装上の大規模検索では、こうした近傍探索をFAISSのような近似最近傍ライブラリで高速化します。

検索で似た過去窓の「続き」が $k$ 本そろいました。これらをそのまま平均してベース予測に足すこともできますが、そこには落とし穴があります。検索結果には「すごく似た当たり」もあれば「ややズレた外れ」も混ざるからです。これを賢く扱うのが、次のARMです。

ARM — 検索結果をMixture-of-Experts的に統合する

ARM(Adaptive Retrieval Mixer)は、$k$ 本の検索結果とバックボーンのベース表現を動的に重み付けして混ぜるモジュールです。発想は Mixture-of-Experts(MoE) に近く、「どの情報源をどれだけ信じるか」をゲートが入力ごとに決めます。検索結果という”専門家”たちと、基盤モデルという”本体”の意見を、状況に応じて配合するイメージです。

順に式を追いましょう。まず、検索してきた各続き $y_i$ を、学習可能な射影器 $f_{\text{MLP}}$ で $d$ 次元ベクトルに埋め込みます。

$$ \begin{equation} \hat{e}_i = f_{\text{MLP}}(y_i), \quad i = 1, 2, \dots, k \end{equation} $$

これらを積み上げて検索表現の行列 $E_{\text{ret}} \in \mathbb{R}^{k \times d}$ を作り、バックボーンが出したクエリ表現 $\hat{e}_q \in \mathbb{R}^{1 \times d}$ と連結します。

$$ \begin{equation} E_{\text{concat}} = [\hat{e}_q;\ E_{\text{ret}}] \in \mathbb{R}^{(k+1) \times d} \end{equation} $$

次に、この連結表現に残差つきの多頭注意(Multi-Head Attention)をかけ、$k+1$ 個の表現どうしの相互作用を学ばせます。

$$ \begin{equation} E_{\text{att}} = \mathrm{MHA}(E_{\text{concat}}) + E_{\text{concat}} \end{equation} $$

注意機構を入れる意味はこうです。「ベース予測と各検索結果が互いをどう補い合うか」を、固定の足し算でなく文脈依存で学べるようにするためです。さらに、ドロップアウトつきの順伝播ネットワーク(FFN)と残差で表現を整えます。

$$ \begin{equation} E_{\text{ffn}} = \mathrm{Dropout}(\mathrm{FFN}(E_{\text{att}})) + E_{\text{att}} \end{equation} $$

ここからが混合(Mixing)の核心です。スコアリングネットワークが、$k+1$ 個の表現それぞれに重要度 $\alpha$ を割り振ります。

$$ \begin{equation} \alpha = \mathrm{Softmax}(W_g E_{\text{ffn}} + b_g) \end{equation} $$

$W_g, b_g$ は学習パラメータ、$\alpha \in \mathbb{R}^{(k+1) \times 1}$ は正規化された注意重みです。Softmaxを通すので、重みは足して1になります。これが「どの検索結果をどれだけ信じるか」のゲートそのものです。最後に、重み付き和を取り、さらにスキップ接続でバックボーンの予測 $\hat{e}_q$ を保ちます。

$$ \begin{equation} e_{\text{final}} = \hat{e}_q + \sum_{i=1}^{k+1} \alpha_i E_{\text{ffn}, i} \end{equation} $$

スキップ接続 $\hat{e}_q$ を明示的に足しているのが効きどころです。検索結果がどれも微妙なときでも、基盤モデル本来の予測力が消えないよう守っているのです。検索が当たれば上乗せ、外れてもベースに戻れる——この安全弁が、ゼロショットでの安定性を支えます。仕上げに、統合表現をバックボーンの射影層 $f_{\text{proj}}$ に通して最終予測を得ます。

$$ \begin{equation} \hat{y}_q = f_{\text{proj}}(e_{\text{final}}) \end{equation} $$

学習はバックボーンと同じ目的関数(Chronos-Boltなら分位点回帰損失)で、ARMと射影器のパラメータだけを更新します。バックボーンとエンコーダは凍結のままです。これで「検索情報を動的に混ぜる小さな部品」を、巨大モデルに後付けで載せられます。

テンソルの形で追う

式(5)〜(12)を「形(shape)」で追うと、ARMが何をしているかがはっきりします。検索した $k$ 本の未来区間 $y_i \in \mathbb{R}^L$ をそれぞれ $f_{\text{MLP}}$ で $d$ 次元に潰し、積み上げて $E_{\text{ret}} \in \mathbb{R}^{k \times d}$。バックボーンのクエリ表現 $\hat{e}_q \in \mathbb{R}^{1 \times d}$ を先頭に連結して $E_{\text{concat}} \in \mathbb{R}^{(k+1) \times d}$。これを「$k+1$ 個のトークン列」とみなして多頭注意に通すので、MHAはベース予測と各検索結果を相互参照させることになります(残差で形は $\mathbb{R}^{(k+1) \times d}$ のまま)。FFN+ドロップアウト+残差で整えた後、スコアリングネット $W_g \in \mathbb{R}^{d \times 1}$ が各トークンに1スカラーを割り当て、Softmaxで $\alpha \in \mathbb{R}^{(k+1) \times 1}$。最後の式(11)は重み付き和 $\sum_i \alpha_i E_{\text{ffn},i} \in \mathbb{R}^{1 \times d}$ に、スキップ接続として $\hat{e}_q$ をもう一度足し込み、$d$ 次元の統合表現 $e_{\text{final}}$ を作ります。これを射影層 $f_{\text{proj}}$ に通すと長さ $L$ の予測 $\hat{y}_q$ が出ます。重要なのは、検索結果が増えても出力の次元 $d$ は不変で、可変個の専門家(検索結果)を固定サイズの表現へ畳み込んでいる点です。これがMoEの「ゲートで専門家を配合する」発想と一致します。

検索を”前段で混ぜる”設計の効率

検索拡張の先行研究 RAF(Retrieval-Augmented Forecasting)は、検索したパターンを生の時系列としてモデル入力に連結する方式でした。これは検索のたびにバックボーンへ長い入力を流し直す必要があり、論文の計測ではETThで検索3290 ms/iter+順伝播184 ms/iter=計3474 ms/iterもかかります。対してTS-RAGは、知識ベースに埋め込みを事前計算して貯めておく(式1の三つ組)ため検索が9.2 ms/iter、ARMによる統合(順伝播)が0.44 ms/iter、計9.62 ms/iter——約360倍速いうえ、平均MSEも0.2318→0.1940とRAFを上回ります(論文Table 4)。「埋め込みを貯める」「未来区間だけを軽い射影器で混ぜる」という二つの設計が、精度と速度を同時に押し上げているわけです。

数式が続いたので、ここで一度「この仕組みが本当に予測を良くするのか」を、手を動かして確かめてみましょう。

Pythonで機構を実証する — 簡易デモ

ここからは、TS-RAGの機構そのものを最小限のコードで再現します。論文のような巨大な基盤モデルや280万ペアの知識ベースは使いません。代わりに、合成した周期+レジーム切替の時系列で「ベース予測 vs 検索拡張した予測」を比べ、誤差が下がることを示します。論文のSOTA性能の再現ではなく、”似た過去の続きを検索して混ぜると効く”という核の確認だと理解してください。

設定の要点はこうです。知識ベースは「きれいな過去のお手本」(低ノイズ)、テスト系列は観測ノイズが大きい、とします。ベース予測(直近1周期のコピー)は1周期ぶんのノイズをそのまま引きずりますが、検索拡張は多数の似た続きを平均してノイズを打ち消せる——これがデモで効く理屈です。

まずデータ生成と「時系列エンコーダ」(形状の指紋を作る簡易版)を定義します。

import numpy as np

PERIOD = 24
REGIMES = [(1.0,1.0,0.0),(1.8,1.0,0.3),(0.6,2.0,-0.2),(1.3,1.0,0.6)]  # (振幅,周波数,トレンド)

def make_series(n_cycles, seed, noise):
    rng = np.random.default_rng(seed); xs = []
    for _ in range(n_cycles):
        amp, freq, trend = REGIMES[rng.integers(len(REGIMES))]
        t = np.arange(PERIOD)
        seg = amp*np.sin(2*np.pi*freq*t/PERIOD) + 0.4*amp*np.sin(2*np.pi*2*freq*t/PERIOD) + trend*(t/PERIOD)
        seg = seg + noise*rng.standard_normal(PERIOD)
        xs.append(seg)
    return np.concatenate(xs)

kb_series   = make_series(400, seed=0,     noise=0.08)   # 知識ベース(きれいなお手本)
test_series = make_series(60,  seed=12345, noise=0.55)   # テスト(観測ノイズ大)

def embed(w):
    """簡易エンコーダ: 平均/標準偏差で正規化→低次フーリエ振幅(位相不変)+1階差統計。"""
    z = (w - w.mean())/(w.std()+1e-8); t = np.arange(len(z)); feats = []
    for k in range(1,6):
        s = np.sin(2*np.pi*k*t/PERIOD); c = np.cos(2*np.pi*k*t/PERIOD)
        feats.append(np.sqrt((z@s)**2 + (z@c)**2)/len(z))
    d1 = np.diff(z)
    feats += [d1.std(), np.mean(np.abs(d1)), z[-PERIOD:].std()]
    return np.array(feats)

エンコーダは波形を平均・分散で正規化したうえで、低次フーリエ振幅(位相に依らずレジームの形を捉える)と変化量の統計を並べた特徴ベクトルにします。スケールやずれに振り回されず「形が似ているか」で比べるための工夫です。次に、この埋め込みで知識ベースを索引化します。

L_IN, L_OUT = 48, 24   # クエリ窓長, 予測長(=1周期)

KB = []
for i in range(0, len(kb_series)-L_IN-L_OUT, 4):
    ctx = kb_series[i:i+L_IN]; fut = kb_series[i+L_IN:i+L_IN+L_OUT]
    KB.append((embed(ctx), ctx, fut))            # (埋め込み, 過去窓, 続き)
KB_emb = np.stack([e for e,_,_ in KB])
KB_fut = np.stack([f for _,_,f in KB])
KB_ctx = np.stack([c for _,c,_ in KB])
emu, esd = KB_emb.mean(0), KB_emb.std(0)+1e-8    # 次元ごと標準化(距離の公平化)
KB_embn = (KB_emb - emu)/esd

これで知識ベース(過去窓・埋め込み・続き)が三つ組でそろいました。続いて、ベース予測・検索・ARM代理の統合を実装します。

def base_forecast(ctx):
    """ベースTSFM代理: 直近1周期をそのまま次周期にコピーする季節ナイーブ。"""
    last = ctx[-PERIOD:]
    return last - last.mean() + ctx[-PERIOD:].mean()

def retrieve(ctx, k):
    eq = (embed(ctx) - emu)/esd
    d = np.linalg.norm(KB_embn - eq, axis=1)     # 式(3) ユークリッド距離
    idx = np.argsort(d)[:k]                       # 式(4) TopK_min
    return idx, d[idx]

def rag_forecast(ctx, k=8, tau=0.5, lam=0.4):
    """検索した続きを距離ベースのソフトマックス重み(ARMのゲート代理)で統合し、
    ベース予測とスキップ接続で混合する。"""
    base = base_forecast(ctx)
    idx, dist = retrieve(ctx, k)
    w = np.exp(-dist/(dist.mean()+1e-8)/tau); w /= w.sum()   # 式(9) α=Softmax
    level = ctx[-PERIOD:].mean()
    retr = np.zeros(L_OUT)
    for wi, j in zip(w, idx):
        retr += wi*(KB_fut[j] - KB_ctx[j][-PERIOD:].mean() + level)  # レベル整合して加重平均
    return (1-lam)*base + lam*retr               # 式(10) スキップ接続つき混合

rag_forecast が ARM の本質を縮約しています。距離に基づくソフトマックスがゲート $\alpha$(式9)に対応し、最後の (1-lam)*base + lam*retr がスキップ接続つきの混合(式10)に対応します。lam は「検索情報をどれだけ信じるか」のつまみです。では、テスト系列全体でベースと比べてみましょう。

def mase(y, yhat, ctx):
    naive = np.mean(np.abs(ctx[PERIOD:] - ctx[:-PERIOD])) + 1e-8
    return np.mean(np.abs(y - yhat))/naive

bm, rm, bmase, rmase = [], [], [], []
for i in range(0, len(test_series)-L_IN-L_OUT, 6):
    ctx = test_series[i:i+L_IN]; y = test_series[i+L_IN:i+L_IN+L_OUT]
    yb = base_forecast(ctx); yr = rag_forecast(ctx, k=8)
    bm.append(np.mean((y-yb)**2)); rm.append(np.mean((y-yr)**2))
    bmase.append(mase(y,yb,ctx));  rmase.append(mase(y,yr,ctx))
print(f"MSE  ベース={np.mean(bm):.4f}  TS-RAG={np.mean(rm):.4f}  改善={(1-np.mean(rm)/np.mean(bm))*100:.1f}%")
print(f"MASE ベース={np.mean(bmase):.4f}  TS-RAG={np.mean(rmase):.4f}  改善={(1-np.mean(rmase)/np.mean(bmase))*100:.1f}%")
# => MSE  ベース=1.0358  TS-RAG=0.8667  改善=16.3%
# => MASE ベース=1.0349  TS-RAG=0.9550  改善=7.7%

出力のとおり、検索拡張で MSEが16.3%、MASEが7.7%下がりました。ノイズの大きいテスト系列に対し、ベースは1周期ぶんのノイズをそのまま予測に持ち込むのに対し、TS-RAGは知識ベースのきれいな似た続きを平均してノイズを打ち消すため、誤差が小さくなります。これがRAGの「実例で補強する」効果の最小再現です。1サンプルの予測を可視化すると、効き方がもっと直感的に見えます。

ベース予測とTS-RAG補正後の予測比較

縦の点線が予測の開始時刻です。青の破線(ベース)は真値(黒)の周りでギザギザに暴れていますが、赤(TS-RAG)はなめらかに真値を追えています。MSEもベース1.128に対しTS-RAG 0.682と、この窓では約4割改善しました。検索で引いた似た区間の続きが、ノイズを均した”お手本”として効いているのが分かります。

では、検索がどんなふうに「似た過去窓」を引いているのか、埋め込み空間を覗いてみましょう。

検索の中身を見る — 埋め込み空間と近傍

検索が機能する条件は「似た形状の窓が埋め込み空間で近くに集まっている」ことです。知識ベースの埋め込みを2次元に主成分射影し、あるクエリ窓とその近傍を重ねて描きます。

埋め込み空間でのクエリ窓と近傍検索

色は各窓の主成分振幅(レジームの目安)です。似た振幅・形状の窓がまとまって雲を作っているのが分かります。赤い星がクエリ窓、赤い丸が検索された上位 $k$ 件で、いずれもクエリのすぐ近くに位置しています。埋め込みの近さが「波形の似ている度合い」に対応しているからこそ、単純なユークリッド距離(式3)で良い近傍が引けるわけです。

検索された近傍の「過去窓そのもの」をクエリ窓と重ねると、形がそろっていることがさらにはっきりします。

検索された近傍の過去窓がクエリと一致

太い黒線がクエリ窓、細い緑線が検索された近傍の過去窓(レベルを合わせて重ねたもの)です。山と谷の位置がよく一致しています。出だしがこれだけ似ていれば、その「続き」も似た振る舞いをすると期待できる——これがTS-RAGの予測補正を支える前提です。

この「形がそろう」現象は、論文の実データ(Weather)でも同じく確認できます。論文は検索の中身を可視化し、解釈可能性の根拠として提示しています。

論文Figure4 Weatherでの検索可視化:クエリと検索された上位3系列の一致

出典: Ning et al., “TS-RAG”, NeurIPS 2025, Fig.4

上段がクエリ系列(青の実線が入力 $x$、青の破線が真の未来 $y$)、下段が知識ベースから引いた Top-1/2/3 の検索系列です。下段の3本がクエリの周期と山谷をきれいになぞっており、トレンドも周期性も一致していることが見て取れます。TS-RAGの検索が単なる数値の近さでなく「波形の構造(周期・トレンド)が似た区間」を選べていることの、論文側の証拠です。しかも検索結果が予測の根拠として人間に提示できる——どの過去パターンを参照したかが見えるので、凍結TSFM単体では得られない解釈可能性が手に入ります。次は「検索を何本引くか」が結果にどう効くかを見ます。

検索本数 k と統合方式の効果

検索する近傍の数 $k$ は重要なつまみです。少なすぎると、たまたま引いた1本のクセに引っ張られて不安定になります。多すぎると、あまり似ていない窓まで混ざってノイズになりかねません。デモで $k$ を振ってみます。

検索本数kに対する誤差の感度

このデモ設定では、$k=1$ のとき誤差が最も大きく(0.94)、$k$ を増やすほど平均化が効いて誤差が下がっていきます。青の破線(ベース)を常に下回っており、検索拡張が一貫して効いていることが分かります。ただし「多ければ良い」はデータ次第で、論文でも知識ベースの構成と $k$ の選び方が性能に効くと報告されています。次に、引いた情報をどう混ぜるかが効くことを確かめます。

論文のアブレーションでは、ARM(注意+ゲート)を、単純に固定重みで足すゲート融合に置き換えると性能が落ちると報告されています。デモでも「RAGなし」「単純平均混合(固定重み)」「ゲート混合(ARM代理)」を比べます。

統合方式の比較ARM>単純混合>RAGなし

RAGなし(1.036)→単純平均混合(0.888)→ゲート混合(0.867)の順に誤差が下がります。検索情報を入れるだけでも効きますが、距離に応じて重みを変えるゲートがさらに上積みを生みます。論文の「ARM > 単純ゲート > RAGなし」という序列が、最小デモでも再現できました。

この序列は、論文のアブレーション(Figure 3)が示すものと同じです。

論文Figure3 アブレーション:融合方式(左)と知識ベース構成(右)

出典: Ning et al., “TS-RAG”, NeurIPS 2025, Fig.3

左の2連バー(Avg. MSE / Avg. MAE)が融合方式の比較です。W/O RAG(赤、バックボーン素のまま)→ W/Gate(橙、検索結果を固定重みで線形結合する単純ゲート)→ W/ARM(青、注意+残差+ゲートのARM)の順に、平均MSEもMAEも下がっています。単純ゲートでもベースは上回りますが、ARMの「注意機構で検索結果どうし・ベースとの相互作用を文脈依存に学ぶ」設計が、さらに低い誤差を生むことが分かります。右側は知識ベースの作り方(Baseline=RAGなし/Cross-Domain/Distribution-Shift/Multi-Domain/In-Domain)の比較で、どの構成でもベースを下回り、なかでも In-Domain(青、予測対象と同じデータから知識ベースを作る)が最小になります。検索拡張は頑健に効き、知識ベースを予測対象に寄せるほど効きが増す、という直感どおりの結果です。そのゲートが実際にどう重みを割り振っているかも、デモで見てみましょう。

ARMのゲート重み

近い(関連の高い)窓ほど大きな重みが割り当てられています。これがARMの式(9)のSoftmaxゲートの働きで、無関係に近い検索結果の影響を自動的に抑える安全装置になっています。最後に、TS-RAGが本領を発揮する「分布シフト」の場面を確かめます。

分布シフトに強い理由

TS-RAGの売りは、非定常な系列・分布シフトに強いことでした。レジームが途中で切り替わる系列を作り、切替の前後で誤差がどう動くかを見ます。

分布シフト前後の予測誤差

灰色の破線がレジーム切替点です。切替の直後はどちらも誤差が跳ねますが、TS-RAG(赤)はベース(青)より低く抑えられ、立ち直りも速い傾向です。新しいレジームに似た窓を知識ベースから引けるので、モデルを再学習しなくても「新しい状況のお手本」を取り込めるからです。重みが凍結された基盤モデル単体では難しい、動的な適応がここで効いています。

機構の確認ができたところで、論文が実データで何を達成したのかを整理しましょう。

評価データセットと主要結果(サーベイ)

論文は、ドメインの異なる7つの公開ベンチマークでゼロショット予測を評価します。ETTh1・ETTh2(変圧器の温度)、ETTm1・ETTm2(同・分単位)、Weather(気象)、Electricity(電力)、Exchange Rate(為替)です。文脈長512、予測長64を基本に、データ分割はETTが6:2:2、その他が7:1:2。評価指標はMSEMAE(小さいほど良い)。比較対象はTTM・TimesFM・Moirai・Chronos・Chronos-Bolt・MOMENT・Time-MoEといった主要TSFMで、TS-RAGはバックボーンにChronos-Boltを採用します。知識ベースとARMの事前学習にはChronos事前学習データから一様サンプリングした多ドメインデータ(知識ベース約280万ペア・事前学習約2600万ペア)を使い、評価対象のデータは一切学習に使わないゼロショット設定です。

論文 Table 1 が主要結果です。まず数値表そのものを引用します。

論文Table1 ゼロショット予測のMSE/MAE主要結果

出典: Ning et al., “TS-RAG”, NeurIPS 2025, Table 1(太字=最良、下線=次点)

左端の TS-RAG$_{\text{Chronos-bolt}}$ 列が、全7データセットでMSE・MAEともに最小(太字)を取っています。バックボーンのChronos-Bolt($\text{Chronos-bolt}_B$)が次点(下線)に並ぶことが多く、検索拡張が「すでに強いバックボーン」をさらに一段持ち上げている構図が読み取れます。とりわけExchange RateはMSE 0.0673→0.0627、ETTm1は0.3109→0.2906と、もとから強い系列でも上積みが効いています。同じ数値を、改善率が見やすいよう自作の棒グラフにもまとめました。

論文Table1ゼロショット予測MSE比較

TS-RAG(赤)が全データセットで最小のMSEを達成しています。バックボーンのChronos-Bolt単体(青)と比べると、平均でMSEを3.54%、MAEを1.43%削減。とりわけ Exchange Rate では、Chronos-Boltがすでに強い(MSE 0.0673)にもかかわらず、TS-RAGがさらに 6.84% 削減して0.0627にしています(これが要旨の「最大+6.84%」の正体で、為替データでのバックボーン比MSE改善率です)。検索拡張が、よく最適化された強いバックボーンの予測すら磨き上げられることを示す結果です。

さらに、長いホライズンで予測を伸ばす(rolling予測)ほど効果が大きくなる点も重要です。論文 Table 2 を見てみましょう。

論文Table2長期予測の改善

ETTh1・Exchangeのどちらも、予測ホライズンが96→192→336→720と伸びるほど、RAGなし(青の破線)とTS-RAG(赤)の差が開いていきます。rolling予測は1ステップごとの誤差が累積しやすいのですが、各ステップで「次の区間の続き」を検索して補強することで、誤差の累積を緩和できるためです。Exchangeのホライズン720では、0.8100→0.6968と大きく改善しています。

論文はこのほか、知識ベースの作り方(in-domain/分布シフト/クロスドメイン/マルチドメイン)による違いも調べ、どの構成でもベースを上回るがin-domainの知識ベースが最も低い誤差になると報告しています。検索拡張が効くこと自体は頑健で、知識ベースの中身を予測対象に寄せるほど効きが増す、という直感どおりの結果です。

まとめ

本記事では、検索拡張で時系列基盤モデルのゼロショット予測を強くするTS-RAG(NeurIPS 2025)を解説しました。

  • 核心: クエリ窓を埋め込み、知識ベースから意味的に近い過去窓を近傍検索し、その続きを予測のヒントに使う。モデルの重みは凍結したまま、外部知識を動的に取り込む
  • ARM(Adaptive Retrieval Mixer): 検索した続きとバックボーン表現を、多頭注意+残差+Softmaxゲート+スキップ接続でMixture-of-Experts的に統合(式5〜12)。$(k+1)\times d$ のトークン列を固定 $d$ 次元へ畳み込む。学習はARMと小さな射影器だけ
  • 効率: 埋め込みを知識ベースに事前計算して貯め、検索した未来区間だけを軽い射影器で混ぜるため、生時系列を入力連結する先行研究RAF(計3474 ms/iter)に対し計9.62 ms/iterと約360倍速く、精度もRAF超え(論文Table 4)
  • 結果: ETT/Weather/Electricity/Exchangeでバックボーン比 平均MSE −3.54%、為替で最大 −6.84%。長いホライズンほど効果が大きく、誤差累積を緩和
  • デモ: 合成データでベース vs 検索拡張を比較し、MSE −16.3%・MASE −7.7%、ARM > 単純混合 > RAGなしの序列を機構として確認(SOTA再現ではない)

検索拡張は、巨大モデルを再学習せずに「いまのデータへ適応」させる軽量で実用的な道具です。知識ベースを差し替えるだけでドメインを移せる柔軟さは、現場運用との相性が良い設計といえます。

次のステップとして、検索パートの実装(ベクトル検索・埋め込み)を深掘りすると、TS-RAGをさらに具体的にイメージできます。