iTransformerとは
Transformerベースの多変量時系列予測モデルの多くは、各時刻の多変量をトークンとみなします。一方iTransformer[1]は、各変数の時系列をトークンとみなすことで変数同士の相関をより良く捉え、予測精度の向上を図りました。大雑把に言うと入出力のテンソルを転置し、vanilla Transformerからモジュールをいくつか引算しただけです。vanilla Transformerとの実装上の相違点を次の表にまとめました。
| 相違点 | vanilla Transformer | iTransformer |
|---|---|---|
| 入出力の形状 | (バッチサイズ, 時系列の長さ, 変数の数) | (バッチサイズ, 変数の数, 時系列の長さ) |
| embedding層の出力の形状 | (バッチサイズ, 時系列の長さ, d_model) | (バッチサイズ, 変数の数, d_model) |
| positional encoding | あり | なし |
| look ahead masks | あり | なし |
| デコーダ | あり | なし |
| 推論 | 次の時刻の多変量を逐次予測 | 時系列を一度に予測 |
所感
従来のTransformerベースの多変量時系列予測モデルは各変数のタイムステップの大きさや時系列の長さが等しいことを前提としていました。しかし、現実にはすべての変数を同じタイムステップで取得できるとは限りませんし、変数ごとに時系列の長さを変えて学習させたいという要望もあるでしょう。各変数のタイムステップの大きさや時系列の長さが異なる場合もアーキテクチャを変更せずに実装できるという点でiTransformerに期待しています。
TensorFlowによる実装
最後に、TensorFlowの公式チュートリアルにあるvanilla Transformerのソースコード[2]をお手本にiTransformerを実装したのでご紹介します。
import tensorflow as tf
def scaled_dot_product_attention(q, k, v):
"""アテンションの重みの計算
q, k, vは最初の次元が一致していること
k, vは最後から2番めの次元が一致していること
引数：
q: query shape == (..., variates_num_q, depth)
k: key shape == (..., variates_num_k, depth)
v: value shape == (..., variates_num_v, depth_v)
戻り値：
出力
"""
matmul_qk = tf.matmul(
q, k, transpose_b=True
) # (..., variates_num_q, variates_num_k)
# matmul_qkをスケール
dk = tf.cast(tf.shape(k)[-1], tf.float32)
scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)
# softmax は最後の軸(variates_num_k)について
# 合計が1となるように正規化
attention_weights = tf.nn.softmax(
scaled_attention_logits, axis=-1
) # (..., variates_num_q, variates_num_k)
output = tf.matmul(attention_weights, v) # (..., variates_num_q, depth_v)
return output
class MultiHeadAttention(tf.keras.layers.Layer):
def __init__(self, d_model, num_heads):
super().__init__()
self.num_heads = num_heads
self.d_model = d_model
assert d_model % self.num_heads == 0
self.depth = d_model // self.num_heads
self.wq = tf.keras.layers.Dense(d_model)
self.wk = tf.keras.layers.Dense(d_model)
self.wv = tf.keras.layers.Dense(d_model)
self.dense = tf.keras.layers.Dense(d_model)
def split_heads(self, x, batch_size):
"""最後の次元を(num_heads, depth)に分割。
結果をshapeが(batch_size, num_heads, variates_num, depth)となるようにリシェイプする。
"""
x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
return tf.transpose(x, perm=[0, 2, 1, 3])
def call(self, v, k, q):
batch_size = tf.shape(q)[0]
q = self.wq(q) # (batch_size, variates_num, d_model)
k = self.wk(k) # (batch_size, variates_num, d_model)
v = self.wv(v) # (batch_size, variates_num, d_model)
q = self.split_heads(
q, batch_size
) # (batch_size, num_heads, variates_num_q, depth)
k = self.split_heads(
k, batch_size
) # (batch_size, num_heads, variates_num_k, depth)
v = self.split_heads(
v, batch_size
) # (batch_size, num_heads, variates_num_v, depth)
# scaled_attention.shape == (batch_size, num_heads, variates_num_q, depth)
scaled_attention = scaled_dot_product_attention(q, k, v)
scaled_attention = tf.transpose(
scaled_attention, perm=[0, 2, 1, 3]
) # (batch_size, variates_num_q, num_heads, depth)
concat_attention = tf.reshape(
scaled_attention, (batch_size, -1, self.d_model)
) # (batch_size, variates_num_q, d_model)
output = self.dense(concat_attention) # (batch_size, variates_num_q, d_model)
return output
def point_wise_feed_forward_network(d_model, dff):
return tf.keras.Sequential(
[
tf.keras.layers.Dense(
dff, activation="gelu"
), # (batch_size, variates_num, dff)
tf.keras.layers.Dense(d_model), # (batch_size, variates_num, d_model)
]
)
class EncoderLayer(tf.keras.layers.Layer):
def __init__(self, d_model, num_heads, dff, rate=0.1):
super().__init__()
self.mha = MultiHeadAttention(d_model, num_heads)
self.ffn = point_wise_feed_forward_network(d_model, dff)
self.dropout1 = tf.keras.layers.Dropout(rate)
self.dropout2 = tf.keras.layers.Dropout(rate)
self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
def call(self, x, training):
attn_output = self.mha(x, x, x) # (batch_size, variates_num, d_model)
attn_output = self.dropout1(attn_output, training=training)
out1 = self.layernorm1(x + attn_output) # (batch_size, variates_num, d_model)
ffn_output = self.ffn(out1) # (batch_size, variates_num, d_model)
ffn_output = self.dropout2(ffn_output, training=training)
out2 = self.layernorm2(out1 + ffn_output) # (batch_size, variates_num, d_model)
return out2
class Encoder(tf.keras.layers.Layer):
def __init__(
self,
num_layers,
d_model,
num_heads,
dff,
rate=0.1,
):
super().__init__()
self.d_model = d_model
self.num_layers = num_layers
self.embedding = tf.keras.layers.Dense(d_model)
self.dropout = tf.keras.layers.Dropout(rate)
self.enc_layers = [
EncoderLayer(d_model, num_heads, dff, rate) for _ in range(num_layers)
]
def call(self, x, training):
x = self.embedding(x) # (batch_size, variates_num, d_model)
x = self.dropout(x, training=training)
for i in range(self.num_layers):
x = self.enc_layers[i](x, training=training)
return x # (batch_size, variates_num, d_model)
class ITransformer(tf.keras.Model):
def __init__(
self,
num_layers,
d_model,
num_heads,
dff,
tar_seq_len,
rate=0.1,
name="i_transformer",
):
super().__init__()
self.encoder = Encoder(num_layers, d_model, num_heads, dff, rate)
self.final_layer = tf.keras.layers.Dense(tar_seq_len)
def call(self, inp, training):
inp = tf.transpose(
inp, perm=[0, 2, 1]
) # (batch_size, variates_num, inp_seq_len)
enc_output = self.encoder(
inp, training=training
) # (batch_size, variates_num, d_model)
final_output = self.final_layer(
enc_output
) # (batch_size, variates_num, tar_seq_len)
final_output = tf.transpose(
final_output, perm=[0, 2, 1]
) # (batch_size, tar_seq_len, variates_num)
return final_output
参考文献
- Liu, Yong, et al. "itransformer: Inverted transformers are effective for time series forecasting." arXiv preprint arXiv:2310.06625 (2023).
- https://www.tensorflow.org/tutorials/text/transformer?hl=ja