iTransformerとは

Transformerベースの多変量時系列予測モデルの多くは、各時刻の多変量をトークンとみなします。一方iTransformer[1]は、各変数の時系列をトークンとみなすことで変数同士の相関をより良く捉え、予測精度の向上を図りました。大雑把に言うと入出力のテンソルを転置し、vanilla Transformerからモジュールをいくつか引算しただけです。vanilla Transformerとの実装上の相違点を次の表にまとめました。

相違点	vanilla Transformer	iTransformer
入出力の形状	(バッチサイズ, 時系列の長さ, 変数の数)	(バッチサイズ, 変数の数, 時系列の長さ)
embedding層の出力の形状	(バッチサイズ, 時系列の長さ, d_model)	(バッチサイズ, 変数の数, d_model)
positional encoding	あり	なし
look ahead masks	あり	なし
デコーダ	あり	なし
推論	次の時刻の多変量を逐次予測	時系列を一度に予測

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