自然言語処理(NLP)がここ数年で飛躍的に進化した背景には、テキストを離散トークンに変換するという基本的なアイデアがあります。「The cat sat on the mat」という文は [464, 3797, 3332, 319, 278, 15082] のようなトークンIDの列に変換され、Transformerはこのトークン列を操作することで言語を理解します。BM25やTF-IDFといった古典的な検索技術も、離散トークンの頻度統計に基づいています。
では、時系列データにも同じアプローチを適用できるでしょうか? 衛星のテレメトリデータ — 温度、電圧、姿勢角速度といった連続値の系列 — を「離散トークンの列」に変換できれば、NLPで蓄積された膨大な技術資産をそのまま時系列に転用できます。BM25で「似たテレメトリパターン」を検索したり、言語モデルで時系列の次のパターンを予測したりすることが可能になります。
しかし、連続値を離散化するには根本的な問題があります。単純なビニング(値域を等間隔に分割)では、データの構造に無関係な量子化が行われるため、重要なパターンが失われます。必要なのは、「データ自身が最適な辞書を学習する」メカニズムです。
TOTEM(Talukder et al., TMLR 2024)は、VQ-VAE(Vector Quantized Variational Autoencoder)を用いてこの課題に取り組んだ手法です。時系列をパッチに分割し、各パッチをVQ-VAEで離散トークンに変換します。学習されたコードブック(離散ベクトルの辞書)は、データに最適化された「テレメトリ辞書」として機能します。
TOTEMを理解することは、以下のような場面で直接役立ちます。
- テレメトリ検索の革新: 時系列を離散トークンに変換することで、BM25やTF-IDFといったNLPの検索技術をそのまま時系列検索に転用できます。ベクトル検索とは異なるアプローチで、解釈可能性の高い検索が可能です
- 言語モデルとの統合: 時系列がトークン列になれば、GPTやLLaMAのような自己回帰言語モデルで時系列の予測・生成が自然に行えます
- データ圧縮: 連続値の時系列をコードブックのインデックスで表現するため、大幅なデータ圧縮が実現します
本記事の内容
- 時系列の離散トークン化はなぜ難しいのか
- VQ-VAEの基礎理論(エンコーダ、コードブック、デコーダ、straight-through estimator)
- VQ-VAEの損失関数の設計と導出
- TOTEMのアーキテクチャ詳細
- コードブック学習と指数移動平均更新
- Pythonによる簡易VQ-VAE実装と時系列トークン化実験
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。

時系列の離散化はなぜ難しいのか
テキストデータの離散化は自然です。単語(またはサブワード)という人間が定義した単位が存在し、語彙サイズ $V$(通常30,000〜50,000)のトークンIDに一対一対応させるだけです。しかし、時系列データには「自然なトークン」が存在しません。
単純なビニングの限界
最も素朴な離散化は、値域を $B$ 個の等間隔ビンに分割し、各時刻の値を対応するビンIDに変換する方法です。Chronos(Amazon, 2024)がこのアプローチを採用しました。
$$ \text{token}(x_t) = \left\lfloor \frac{x_t – x_{\min}}{x_{\max} – x_{\min}} \cdot B \right\rfloor $$
この方法には2つの根本的な問題があります。
量子化誤差: ビン幅 $\delta = (x_{\max} – x_{\min}) / B$ の半分が最大誤差です。$B = 4096$ でも、値域が広い場合(例: 電力需要の0〜50,000 MW)では $\delta \approx 12$ MWとなり、微細な変動パターンが消失します。
時間構造の無視: 各時刻を独立にビン分けするため、「上昇トレンドの途中」と「下降トレンドの途中」で同じ値を取る時刻が同じトークンに割り当てられます。時間的な文脈が完全に失われます。
パッチ + 学習型離散化
これらの問題を解決するアイデアは、パッチ化と学習型コードブックの組み合わせです。時系列を固定長のパッチ(部分系列)に分割し、各パッチを「形状が似ているもの同士」でクラスタリングしてトークンIDを割り当てます。ここでVQ-VAEが登場します。VQ-VAEは、パッチの圧縮表現を学習しながら、同時にコードブック(プロトタイプベクトルの集合)も学習する枠組みです。
時系列の離散化の困難さがわかったところで、VQ-VAEがこの問題をどう解決するか、その基礎理論から見ていきましょう。
VQ-VAEの基礎理論
VAEからVQ-VAEへ
通常のVAE(Variational Autoencoder)は、入力 $\bm{x}$ を連続的な潜在変数 $\bm{z} \in \mathbb{R}^d$ にエンコードし、デコーダで再構成します。潜在空間は連続であり、事前分布 $p(\bm{z}) = \mathcal{N}(\bm{0}, \bm{I})$ との KL ダイバージェンスを正則化項として加えます。
しかし、連続的な潜在空間には「後方崩壊(posterior collapse)」という問題があります。デコーダが十分に強力だと、潜在変数を無視してデコーダだけで入力を再構成しようとし、潜在表現が意味を持たなくなります。
VQ-VAE(van den Oord et al., 2017)は、この問題を潜在空間の離散化で解決します。連続的な潜在ベクトルの代わりに、コードブック $\mathcal{C} = \{\bm{e}_1, \bm{e}_2, \ldots, \bm{e}_K\}$($K$ 個のプロトタイプベクトル、各 $\bm{e}_k \in \mathbb{R}^d$)を用意し、エンコーダ出力を最も近いコードブックベクトルで置き換えます。
$$ \bm{z}_q = \bm{e}_{k^*}, \quad k^* = \arg\min_{k \in \{1, \ldots, K\}} \|\bm{z}_e – \bm{e}_k\|_2 $$
ここで $\bm{z}_e = f_{\text{enc}}(\bm{x})$ はエンコーダの出力、$\bm{z}_q$ は量子化後のベクトルです。入力はコードブックのインデックス $k^*$ で一意に表現され、これが「離散トークン」となります。
イメージとしては、画家が無限の色から選ぶ代わりに、あらかじめ用意された $K$ 色のパレットから最も近い色を選ぶようなものです。パレットの色は固定ではなく、より良い表現が得られるように学習で更新されます。
エンコーダ・コードブック・デコーダの3層構造
VQ-VAEの全体構造は以下の3つのコンポーネントで構成されます。
エンコーダ $f_{\text{enc}}: \mathbb{R}^P \to \mathbb{R}^d$ — 入力パッチ $\bm{x} \in \mathbb{R}^P$(パッチサイズ $P$)を潜在ベクトル $\bm{z}_e \in \mathbb{R}^d$ に変換します。
量子化器(Quantizer) — $\bm{z}_e$ に最も近いコードブックベクトルを選択します。
$$ q(\bm{z}_e) = \bm{e}_{k^*}, \quad k^* = \arg\min_k \|\bm{z}_e – \bm{e}_k\|_2^2 $$
デコーダ $f_{\text{dec}}: \mathbb{R}^d \to \mathbb{R}^P$ — 量子化後のベクトル $\bm{z}_q$ から元のパッチ $\hat{\bm{x}}$ を再構成します。
情報の流れを数式で書くと以下のようになります。
$$ \bm{x} \xrightarrow{f_{\text{enc}}} \bm{z}_e \xrightarrow{q(\cdot)} \bm{z}_q \xrightarrow{f_{\text{dec}}} \hat{\bm{x}} $$
この3層構造でエンコーダとデコーダは微分可能ですが、量子化器の $\arg\min$ 操作は微分不可能です。これがVQ-VAEの学習における最大の課題であり、straight-through estimator で解決します。
VQ-VAEの構造を理解したところで、次にこの微分不可能性をどう乗り越えるか、そして全体の損失関数がどう設計されているかを見ていきます。
VQ-VAEの損失関数
微分不可能性とstraight-through estimator
量子化器 $q(\bm{z}_e) = \bm{e}_{k^*}$ には $\arg\min$ 操作が含まれており、勾配が定義できません。van den Oord et al. (2017) はstraight-through estimator(STE)を用いてこの問題を解決しました。
アイデアはシンプルです。順伝播では量子化後のベクトル $\bm{z}_q$ を使い、逆伝播ではエンコーダ出力 $\bm{z}_e$ にそのまま勾配を流すという近似です。数式で書くと、
$$ \bm{z}_q = \bm{z}_e + \text{sg}[\bm{z}_q – \bm{z}_e] $$
ここで $\text{sg}[\cdot]$ は stop-gradient 操作(順伝播はそのまま通すが、逆伝播では勾配をゼロにする)です。この式を展開すると、
順伝播: $\bm{z}_q = \bm{z}_e + (\bm{z}_q – \bm{z}_e) = \bm{z}_q$(量子化後のベクトル)
逆伝播: $\frac{\partial \bm{z}_q}{\partial \bm{z}_e} = 1 + 0 = 1$(勾配がそのまま通過)
つまり、デコーダからの勾配がエンコーダにそのまま伝わります。これにより、エンコーダは「量子化後のベクトルがデコーダにとって有用になるような潜在表現」を学習できます。
3つの損失項
VQ-VAEの損失関数は3つの項から構成されます。
$$ \mathcal{L} = \underbrace{\|\bm{x} – \hat{\bm{x}}\|_2^2}_{\text{再構成損失}} + \underbrace{\|\text{sg}[\bm{z}_e] – \bm{z}_q\|_2^2}_{\text{コードブック損失}} + \underbrace{\beta \|\bm{z}_e – \text{sg}[\bm{z}_q]\|_2^2}_{\text{コミットメント損失}} $$
各項の役割を詳しく見ていきましょう。
第1項: 再構成損失 $\|\bm{x} – \hat{\bm{x}}\|_2^2$
入力パッチ $\bm{x}$ と再構成パッチ $\hat{\bm{x}} = f_{\text{dec}}(\bm{z}_q)$ の二乗誤差です。STEにより、この損失の勾配はエンコーダとデコーダの両方に伝播します。この項はモデル全体の表現力を駆動します。
第2項: コードブック損失 $\|\text{sg}[\bm{z}_e] – \bm{z}_q\|_2^2$
コードブックベクトルをエンコーダ出力に近づける項です。$\text{sg}[\bm{z}_e]$ はエンコーダ出力を定数として扱うため、この項の勾配はコードブックベクトルのみに作用します。直感的には、「コードブックがエンコーダ出力の方に歩み寄る」ための力です。
第3項: コミットメント損失 $\beta \|\bm{z}_e – \text{sg}[\bm{z}_q]\|_2^2$
エンコーダ出力をコードブックベクトルに近づける項です。$\text{sg}[\bm{z}_q]$ はコードブックベクトルを定数として扱うため、この項の勾配はエンコーダのみに作用します。直感的には、「エンコーダがコードブックの方に歩み寄る」ための力です。$\beta$ は通常0.25に設定されます。
第2項と第3項が対称的な役割を果たしている点がポイントです。コードブックとエンコーダが互いに歩み寄ることで、量子化誤差 $\|\bm{z}_e – \bm{z}_q\|$ が小さくなり、STEの近似精度が向上します。
なぜKLダイバージェンス項が不要なのか
通常のVAEでは、潜在変数の事後分布 $q(\bm{z}|\bm{x})$ と事前分布 $p(\bm{z})$ のKLダイバージェンスを正則化項として加えます。VQ-VAEではこれが不要です。なぜなら、潜在空間が離散化されており、事後分布は $q(k|\bm{x}) = \mathbb{1}[k = k^*]$(最も近いコードのみ確率1)というデルタ分布になります。事前分布をカテゴリカル分布の一様分布 $p(k) = 1/K$ とすると、
$$ \text{KL}[q(k|\bm{x}) \| p(k)] = \log K $$
これは全ての入力に対して定数なので、最適化に影響しません。VQ-VAEは明示的にKLダイバージェンスを計算する必要がなく、損失関数が簡潔になるのです。
損失関数の全体像を理解したところで、次にTOTEMがVQ-VAEを時系列にどう適用するかを見ていきます。
TOTEMのアーキテクチャ
時系列のパッチ化
TOTEMは入力時系列をまずパッチに分割します。長さ $L$ の単変量時系列 $\bm{x} = (x_1, x_2, \ldots, x_L)$ をパッチサイズ $P$、ストライド $S$ で分割すると、$N = \lfloor (L – P) / S \rfloor + 1$ 個のパッチが得られます。
$$ \bm{p}_i = (x_{(i-1)S+1}, x_{(i-1)S+2}, \ldots, x_{(i-1)S+P}), \quad i = 1, \ldots, N $$
各パッチ $\bm{p}_i \in \mathbb{R}^P$ が VQ-VAE の入力単位となります。衛星テレメトリの場合、サンプリングレートが1Hzであれば $P = 60$(1分間のパッチ)や $P = 300$(5分間のパッチ)が典型的な設定です。
パッチ化の重要な効果は、時間的な局所構造を保持したまま離散化できる点です。個々の値ではなく、波形の「形状」単位でトークン化されるため、上昇トレンド、周期的振動、異常パルスといったパターンがコードブックのエントリとして学習されます。
エンコーダとデコーダの設計
TOTEMのエンコーダとデコーダには、1次元の畳み込みネットワーク(Conv1D)が使われます。
エンコーダ: 複数の Conv1D 層(カーネルサイズ4、ストライド2)でパッチを圧縮し、最終的に $d$ 次元の潜在ベクトルを出力します。各層の後にバッチ正規化とReLU活性化が適用されます。
$$ \bm{z}_e^{(i)} = f_{\text{enc}}(\bm{p}_i) \in \mathbb{R}^d $$
デコーダ: エンコーダの逆構造(転置畳み込み)で、量子化後のベクトルからパッチを再構成します。
$$ \hat{\bm{p}}_i = f_{\text{dec}}(\bm{z}_q^{(i)}) \in \mathbb{R}^P $$
トークン列の生成
$N$ 個のパッチそれぞれがコードブックのインデックスに変換されるため、時系列全体は長さ $N$ のトークン列で表現されます。
$$ \text{tokens}(\bm{x}) = (k_1^*, k_2^*, \ldots, k_N^*) $$
ここで $k_i^* = \arg\min_k \|\bm{z}_e^{(i)} – \bm{e}_k\|_2^2$ です。コードブックサイズ $K = 512$ の場合、各トークンは9ビット($\log_2 512 = 9$)で表現でき、これは元のパッチ($P$ 個の32ビット浮動小数点数 = $32P$ ビット)に比べて大幅な圧縮です。パッチサイズ $P = 64$ なら圧縮比は $32 \times 64 / 9 \approx 227$ 倍です。
TOTEMのアーキテクチャが明らかになりました。次は、コードブックの学習を安定させるための重要なテクニックについて見ていきます。
コードブック学習の安定化
指数移動平均(EMA)更新
コードブック損失 $\|\text{sg}[\bm{z}_e] – \bm{z}_q\|_2^2$ を勾配降下法で最適化する方法は有効ですが、実践ではより安定した指数移動平均(EMA)更新が広く使われます。
各コードブックベクトル $\bm{e}_k$ は、学習中にそれに割り当てられたエンコーダ出力の指数移動平均として更新されます。ミニバッチ内で $\bm{e}_k$ に割り当てられたエンコーダ出力の集合を $\mathcal{S}_k = \{\bm{z}_e^{(j)} : k^{*(j)} = k\}$ とすると、
まず割り当て数の移動平均を計算します。
$$ N_k \leftarrow \gamma N_k + (1 – \gamma) |\mathcal{S}_k| $$
次に割り当てベクトルの移動平均を計算します。
$$ \bm{m}_k \leftarrow \gamma \bm{m}_k + (1 – \gamma) \sum_{\bm{z} \in \mathcal{S}_k} \bm{z} $$
最後にコードブックベクトルを更新します。
$$ \bm{e}_k \leftarrow \frac{\bm{m}_k}{N_k} $$
ここで $\gamma = 0.99$ が典型的な減衰率です。この更新は $k$-means の重心更新のオンライン版と解釈できます。
コードブック崩壊の防止
VQ-VAEの学習でよく起こる問題がコードブック崩壊(codebook collapse)です。これは、一部のコードブックベクトルだけが頻繁に使われ、残りのベクトルが一度も割り当てられなくなる現象です。$K = 512$ のコードブックで実際に使われるのが50個だけ、ということが起こり得ます。
崩壊が発生すると、表現力が大幅に低下します。512個のコードで表現可能な $512^N$ 通りのパターンが、50個のコードによる $50^N$ 通りに減ってしまいます。
TOTEMでは以下の対策が講じられています。
ランダム再初期化: 一定ステップ数連続で割り当てがなかったコードブックベクトルを、ランダムに選んだエンコーダ出力で再初期化します。
$$ \bm{e}_k \leftarrow \bm{z}_e^{(\text{random})}, \quad \text{if } |\mathcal{S}_k| = 0 \text{ for } T_{\text{dead}} \text{ steps} $$
コードブック使用率のモニタリング: 学習中にコードブックの使用率($|\{k : |\mathcal{S}_k| > 0\}| / K$)を監視し、過度な崩壊を検出します。
温度付きソフト量子化(学習初期): $\arg\min$ の代わりにソフトマックスベースの確率的割り当てを使うことで、初期段階での多様なコード使用を促進します。
$$ p(k | \bm{z}_e) = \frac{\exp(-\|\bm{z}_e – \bm{e}_k\|_2^2 / \tau)}{\sum_{j=1}^K \exp(-\|\bm{z}_e – \bm{e}_j\|_2^2 / \tau)} $$
温度 $\tau$ を学習の進行とともに下げていくことで、最終的にハードな量子化に収束させます。
コードブック学習のテクニックを理解したところで、離散トークンが得られた後に何ができるのか、応用の可能性について見ていきましょう。
離散トークン化の応用
NLP検索技術の転用
時系列がトークン列に変換されると、テキスト検索で使われるBM25を直接適用できます。
BM25スコアは各トークン(ここではコードブックインデックス)の出現頻度に基づいて定義されます。クエリ時系列のトークン列 $\bm{q} = (q_1, \ldots, q_M)$ と、データベース内の時系列のトークン列 $\bm{d} = (d_1, \ldots, d_N)$ に対して、
$$ \text{BM25}(\bm{q}, \bm{d}) = \sum_{t \in \bm{q}} \text{IDF}(t) \cdot \frac{f(t, \bm{d}) \cdot (k_1 + 1)}{f(t, \bm{d}) + k_1 \cdot \left(1 – b + b \cdot \frac{|\bm{d}|}{d_{\text{avg}}}\right)} $$
ここで $f(t, \bm{d})$ はトークン $t$ のドキュメント $\bm{d}$ 内での出現頻度、$\text{IDF}(t)$ は逆文書頻度です。
これは「このテレメトリパターン(トークン)が、他の時系列と比べてどれだけ特徴的か」を定量化していることになります。珍しいパターン(高IDF値)が一致すれば高スコアとなり、よくあるパターン(低IDF値)の一致は低スコアに抑えられます。
言語モデルによる時系列生成
トークン列が得られれば、GPT型の自己回帰モデルで次のトークンを予測できます。
$$ p(k_{N+1} | k_1, \ldots, k_N) = \text{Transformer}(k_1, \ldots, k_N) $$
予測されたトークン $k_{N+1}$ をコードブックベクトル $\bm{e}_{k_{N+1}}$ に変換し、デコーダを通すことで時系列の予測値が得られます。この方法の利点は、連続値の回帰問題を分類問題に変換できることです。$K$ クラスの分類として定式化されるため、予測の不確実性は $K$ 個のクラスの確率分布として自然に表現されます。
データ圧縮と転送
衛星通信では帯域幅が極めて限られています。32ビット浮動小数点の生テレメトリを送信する代わりに、コードブックインデックス(例: 9ビット)を送信することで、$32P / \log_2 K$ 倍の圧縮が可能です。地上局でコードブックを共有しておけば、受信側でデコーダを使って時系列を再構成できます。
応用の可能性を確認したところで、実際にPythonでVQ-VAEを実装し、時系列をトークン化してみましょう。
PythonによるVQ-VAE実装
簡易VQ-VAEモデル
以下のコードでは、PyTorchを使わずにNumPyのみで簡易VQ-VAEの学習ループを実装します。エンコーダとデコーダは線形変換で近似し、コードブック学習の核心部分に焦点を当てます。
import numpy as np
import matplotlib.pyplot as plt
np.random.seed(42)
# --- 合成時系列データの生成 ---
def generate_patterns(n_samples=500, patch_size=32):
"""5種類の基本パターンを持つ時系列パッチを生成"""
t = np.linspace(0, 2 * np.pi, patch_size)
patterns = []
labels = []
for _ in range(n_samples):
pattern_type = np.random.randint(0, 5)
noise = np.random.randn(patch_size) * 0.1
if pattern_type == 0: # 上昇トレンド
p = np.linspace(-1, 1, patch_size) + noise
elif pattern_type == 1: # 下降トレンド
p = np.linspace(1, -1, patch_size) + noise
elif pattern_type == 2: # 正弦波
freq = np.random.uniform(1, 3)
p = np.sin(freq * t) + noise
elif pattern_type == 3: # パルス(異常)
p = noise.copy()
pulse_pos = np.random.randint(8, 24)
p[pulse_pos:pulse_pos+4] += 3.0
else: # 定常(フラット)
p = np.random.randn(1) * 0.3 + noise
patterns.append(p)
labels.append(pattern_type)
return np.array(patterns), np.array(labels)
patches, true_labels = generate_patterns(n_samples=1000, patch_size=32)
# --- VQ-VAE コンポーネント ---
class SimpleVQVAE:
def __init__(self, input_dim=32, latent_dim=8, codebook_size=16,
beta=0.25, lr=0.01, ema_decay=0.99):
self.latent_dim = latent_dim
self.codebook_size = codebook_size
self.beta = beta
self.lr = lr
self.ema_decay = ema_decay
# エンコーダ(線形変換)
self.W_enc = np.random.randn(input_dim, latent_dim) * 0.1
self.b_enc = np.zeros(latent_dim)
# デコーダ(線形変換)
self.W_dec = np.random.randn(latent_dim, input_dim) * 0.1
self.b_dec = np.zeros(input_dim)
# コードブック
self.codebook = np.random.randn(codebook_size, latent_dim) * 0.5
# EMA用
self.ema_count = np.ones(codebook_size)
self.ema_sum = self.codebook.copy()
def encode(self, x):
return np.tanh(x @ self.W_enc + self.b_enc)
def quantize(self, z_e):
# 各z_eに最も近いコードブックベクトルを選択
dists = np.sum((z_e[:, None, :] - self.codebook[None, :, :]) ** 2, axis=2)
indices = np.argmin(dists, axis=1)
z_q = self.codebook[indices]
return z_q, indices
def decode(self, z_q):
return z_q @ self.W_dec + self.b_dec
def update_codebook_ema(self, z_e, indices):
for k in range(self.codebook_size):
mask = (indices == k)
count = np.sum(mask)
if count > 0:
self.ema_count[k] = self.ema_decay * self.ema_count[k] + (1 - self.ema_decay) * count
self.ema_sum[k] = self.ema_decay * self.ema_sum[k] + (1 - self.ema_decay) * np.sum(z_e[mask], axis=0)
self.codebook[k] = self.ema_sum[k] / self.ema_count[k]
elif np.random.rand() < 0.01: # 未使用コードの再初期化
rand_idx = np.random.randint(len(z_e))
self.codebook[k] = z_e[rand_idx]
self.ema_sum[k] = z_e[rand_idx]
self.ema_count[k] = 1.0
def train_step(self, x_batch):
# 順伝播
z_e = self.encode(x_batch)
z_q, indices = self.quantize(z_e)
# Straight-through estimator: 逆伝播ではz_eの勾配を使う
z_st = z_e + (z_q - z_e) # 順伝播はz_q、逆伝播はz_eとして近似
x_recon = self.decode(z_st)
# 損失計算
recon_loss = np.mean((x_batch - x_recon) ** 2)
commit_loss = np.mean((z_e - z_q) ** 2)
total_loss = recon_loss + self.beta * commit_loss
# コードブックのEMA更新
self.update_codebook_ema(z_e, indices)
# 簡易的な勾配更新(エンコーダとデコーダ)
grad_recon = -2 * (x_batch - x_recon) / len(x_batch)
# デコーダ更新
grad_W_dec = z_st.T @ grad_recon / len(x_batch)
grad_b_dec = np.mean(grad_recon, axis=0)
self.W_dec -= self.lr * grad_W_dec
self.b_dec -= self.lr * grad_b_dec
# エンコーダ更新(STE + コミットメント損失の勾配)
grad_z = grad_recon @ self.W_dec.T
grad_commit = 2 * self.beta * (z_e - z_q) / len(x_batch)
grad_z_total = grad_z + grad_commit
dtanh = 1 - z_e ** 2 # tanh の導関数
grad_z_total *= dtanh
grad_W_enc = x_batch.T @ grad_z_total / len(x_batch)
grad_b_enc = np.mean(grad_z_total, axis=0)
self.W_enc -= self.lr * grad_W_enc
self.b_enc -= self.lr * grad_b_enc
# コードブック使用率
usage = len(np.unique(indices)) / self.codebook_size
return total_loss, recon_loss, commit_loss, usage
def tokenize(self, x):
z_e = self.encode(x)
_, indices = self.quantize(z_e)
return indices
# --- 学習ループ ---
model = SimpleVQVAE(input_dim=32, latent_dim=8, codebook_size=16, beta=0.25, lr=0.005)
n_epochs = 100
batch_size = 64
losses = []
usages = []
for epoch in range(n_epochs):
perm = np.random.permutation(len(patches))
epoch_loss = 0
epoch_usage = 0
n_batches = 0
for i in range(0, len(patches), batch_size):
batch = patches[perm[i:i+batch_size]]
if len(batch) < 2:
continue
loss, rl, cl, usage = model.train_step(batch)
epoch_loss += loss
epoch_usage += usage
n_batches += 1
losses.append(epoch_loss / n_batches)
usages.append(epoch_usage / n_batches)
# --- 可視化: 学習曲線 ---
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
axes[0].plot(losses, color='#00d4ff')
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Total Loss')
axes[0].set_title('VQ-VAE Training Loss')
axes[0].grid(True, alpha=0.3)
axes[1].plot(usages, color='#ff6b6b')
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('Codebook Usage')
axes[1].set_title('Codebook Utilization Rate')
axes[1].set_ylim(0, 1.1)
axes[1].grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('vqvae_training.png', dpi=150, bbox_inches='tight')
plt.show()
上のグラフから、2つの重要な特徴が読み取れます。左のグラフでは、再構成損失が学習の進行とともに単調に減少しており、VQ-VAEが時系列パッチの圧縮表現を効果的に学習していることがわかります。右のグラフはコードブック使用率を示しており、EMA更新とランダム再初期化の効果で、多くのコードブックエントリが活用されていることが確認できます。使用率が1.0に近いほど、コードブック崩壊が防止されていることを意味します。
トークン化の結果を可視化
学習済みモデルで時系列をトークン化し、各トークンがどのようなパターンに対応するかを可視化します。
# --- トークン化と可視化 ---
tokens = model.tokenize(patches)
# 各トークンに割り当てられたパッチの代表例を表示
fig, axes = plt.subplots(4, 4, figsize=(14, 10))
for k in range(16):
ax = axes[k // 4, k % 4]
mask = (tokens == k)
if np.sum(mask) > 0:
assigned = patches[mask]
# 最大10個のパッチを重ねて表示
for j in range(min(10, len(assigned))):
ax.plot(assigned[j], alpha=0.3, color='#00d4ff')
# 平均パッチを太線で表示
ax.plot(np.mean(assigned, axis=0), color='#ff6b6b', linewidth=2)
ax.set_title(f'Token {k} (n={np.sum(mask)})', fontsize=9)
else:
ax.set_title(f'Token {k} (unused)', fontsize=9)
ax.set_xlim(0, 31)
ax.grid(True, alpha=0.2)
ax.tick_params(labelsize=7)
plt.suptitle('Codebook Entries: Each Token Represents a Time Series Pattern', fontsize=12)
plt.tight_layout()
plt.savefig('vqvae_codebook.png', dpi=150, bbox_inches='tight')
plt.show()
この可視化から、各トークンが特定のパターンタイプに対応していることが確認できます。例えば、あるトークンは上昇トレンドのパッチばかりが割り当てられ、別のトークンは正弦波パターンのパッチが集まっています。パルス(異常)パターンは独自のトークンに割り当てられる傾向があり、これはコードブックが「テレメトリ辞書」として機能していることを裏付けます。赤い太線は各トークンの平均パッチ(プロトタイプ波形)を示しており、コードブックベクトルをデコードした結果に対応します。
BM25による時系列検索
離散トークンが得られたので、BM25で時系列検索を行います。
# --- BM25による時系列検索 ---
from collections import Counter
def compute_idf(token_sequences, K):
"""逆文書頻度を計算"""
N = len(token_sequences)
df = np.zeros(K)
for seq in token_sequences:
unique_tokens = set(seq) if hasattr(seq, '__iter__') else {seq}
for t in unique_tokens:
df[t] += 1
idf = np.log((N - df + 0.5) / (df + 0.5) + 1)
return idf
def bm25_score(query_tokens, doc_tokens, idf, k1=1.5, b=0.75, avg_dl=10):
"""BM25スコアを計算"""
dl = len(doc_tokens)
tf = Counter(doc_tokens)
score = 0.0
for t in query_tokens:
f = tf.get(t, 0)
numerator = f * (k1 + 1)
denominator = f + k1 * (1 - b + b * dl / avg_dl)
score += idf[t] * numerator / denominator
return score
# 全パッチをトークン化(個別トークンだが、連続パッチのシーケンスとして扱う)
# 長い時系列を模擬: 10パッチずつまとめてドキュメントとする
n_docs = len(tokens) // 10
token_docs = [tokens[i*10:(i+1)*10].tolist() for i in range(n_docs)]
idf = compute_idf(token_docs, model.codebook_size)
# クエリ: パルスパターンのパッチ
pulse_idx = np.where(true_labels == 3)[0][:10]
query_tokens = model.tokenize(patches[pulse_idx]).tolist()
# 検索
scores = []
for doc in token_docs:
scores.append(bm25_score(query_tokens, doc, idf))
scores = np.array(scores)
# 上位5件を表示
top_k = 5
top_indices = np.argsort(scores)[::-1][:top_k]
fig, axes = plt.subplots(1, top_k, figsize=(16, 3))
for i, idx in enumerate(top_indices):
doc_patches = patches[idx*10:(idx+1)*10]
for p in doc_patches:
axes[i].plot(p, alpha=0.5, color='#00d4ff')
axes[i].set_title(f'Rank {i+1}\nBM25={scores[idx]:.2f}', fontsize=9)
axes[i].grid(True, alpha=0.2)
plt.suptitle('BM25 Search Results: Query = Pulse (Anomaly) Pattern', fontsize=11)
plt.tight_layout()
plt.savefig('vqvae_bm25_search.png', dpi=150, bbox_inches='tight')
plt.show()
print(f"クエリトークン: {query_tokens}")
print(f"IDF値(トークンごと): {np.round(idf, 2)}")
BM25検索の結果から、パルス(異常)パターンを含むドキュメントが上位にランクされていることが確認できます。特にスコアが高いドキュメントほど、クエリと同じトークン(パルスパターンに対応するコードブックインデックス)を多く含んでいます。IDF値を見ると、パルスパターンのトークンは出現頻度が低いため高いIDF値を持ち、BM25スコアに大きく寄与していることがわかります。これはNLPにおける「稀少な単語ほど検索に有用」という原則がそのまま時系列にも適用されていることを示しています。
まとめ
本記事では、TOTEM(TMLR 2024)が提案するVQ-VAEベースの時系列離散トークン化について解説しました。
- VQ-VAEの核心: エンコーダ出力を最も近いコードブックベクトルに量子化する。微分不可能な量子化操作はstraight-through estimatorで対処し、再構成損失・コードブック損失・コミットメント損失の3項で学習する
- コードブックはデータ適応型辞書: 等間隔ビニングとは異なり、データ自身が最適な離散表現を学習する。EMA更新とコードブック崩壊防止策で安定的に学習できる
- 離散トークン化の3つの応用: NLP検索技術(BM25)の転用、言語モデルによる時系列生成、データ圧縮と帯域削減
TOTEMのアプローチは、時系列データとNLP技術の橋渡しとして重要な位置を占めます。特にテレメトリ検索においては、ベクトル検索(FAISS等)とは異なる「トークンベースの検索」という選択肢を提供し、解釈可能性とスケーラビリティの両立を可能にします。
次のステップとして、以下の記事も参考にしてください。