在多变量时序预测任务中，如何利用多头注意力机制计算各变量之间的相关系

Question

在多变量时序预测任务中，如何利用多头注意力机制求得各变量之间的相关系？
例如输入为：（100,4,12）其中12为多个变量，我想计算变量之间的相关性，下面代码好像是计算时间步之间的相关性，求各位指点指点哇

class MultiHeadSelfAttention(Layer):
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadSelfAttention, self).__init__()

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        if embed_dim % num_heads != 0:
            raise ValueError(
                f"embedding dimension = {embed_dim} should be divisible by number of heads = {num_heads}"
            )

        self.projection_dim = embed_dim // num_heads
        self.query_dense = Dense(embed_dim)
        self.key_dense = Dense(embed_dim)
        self.value_dense = Dense(embed_dim)
        self.output_dense = Dense(embed_dim)
        # self.attention=MultiHeadAttention()

    def attention(self, query, key, value):
        score = tf.matmul(query, key, transpose_b=True)
        dim_key = tf.cast(tf.shape(key)[-1], tf.float32)
        scaled_score = score / tf.math.sqrt(dim_key)
        weights = tf.nn.softmax(scaled_score, axis=-1)
        out = tf.matmul(weights, value)
        return out, weights

    def separate_heads(self, x, batch_size):
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.projection_dim))
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, inputs):
        # x.shape = [batch_size, seq_len, embedding_dim]
        batch_size = tf.shape(inputs)[0]

        # Query
        # Dense layer to project and split embed_dim to num_heads*projection_dim
        query = self.query_dense(inputs)  # (batch_size, seq_len, embed_dim)
        query = self.separate_heads(
            query, batch_size
        )  # (batch_size, num_heads, seq_len, projection_dim)

        # Key
        # Dense layer to project and split embed_dim to num_heads*projection_dim
        key = self.key_dense(inputs)  # (batch_size, seq_len, embed_dim)
        key = self.separate_heads(
            key, batch_size
        )  # (batch_size, num_heads, seq_len, projection_dim)

        # Value
        # Dense layer to project and split embed_dim to num_heads*projection_dim
        value = self.value_dense(inputs)  # (batch_size, seq_len, embed_dim)
        value = self.separate_heads(
            value, batch_size
        )  # (batch_size, num_heads, seq_len, projection_dim)

        # Attention
        attention, weigths = self.attention(query, key, value)
        attention = tf.transpose(
            attention, perm=[0, 2, 1, 3]
        )  # (batch_size, seq_len, num_heads, projection_dim)

        # Combine back to num_heads*projection_dim
        concat_attention = tf.reshape(
            attention, (batch_size, -1, self.embed_dim)
        )  # (batch_size, seq_len, embed_dim)

        # Apply a final Dense layer
        out = self.output_dense(
            concat_attention
        )  # (batch_size, seq_len, embed_dim)
        return out

Answer 1

引用new bing作答：
在多变量时序预测任务中，利用多头注意力机制求得各变量之间的相关系数，需要对输入的数据进行一些处理和调整。假设输入数据的维度为(batch_size, seq_len, num_variables)，其中num_variables为变量的数量。具体的步骤如下：

1 将输入的数据进行转置，使得变量的数量成为第二个维度，即新的维度为(batch_size, num_variables, seq_len)。

2 将转置后的数据输入到多头注意力机制中，通过query, key和value的计算，可以得到注意力矩阵，注意力矩阵的维度为(batch_size, num_heads, seq_len, seq_len)，其中seq_len为序列的长度，num_heads为注意力头的数量。

3 将注意力矩阵进行汇总，得到每个时间步与其他时间步的相关系数，具体的计算方法为将每个头的注意力矩阵相加并除以num_heads，得到的维度为(batch_size, seq_len, seq_len)。

4 最后，将相关系数矩阵进行转置，使得变量的数量成为第一个维度，即新的维度为(batch_size, num_variables, num_variables)，即可得到各变量之间的相关系数矩阵。

下面是代码实现：

class MultiHeadAttention(Layer):
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        if embed_dim % num_heads != 0:
            raise ValueError(
                f"embedding dimension = {embed_dim} should be divisible by number of heads = {num_heads}"
            )

        self.projection_dim = embed_dim // num_heads
        self.query_dense = Dense(embed_dim)
        self.key_dense = Dense(embed_dim)
        self.value_dense = Dense(embed_dim)
        self.output_dense = Dense(embed_dim)

    def attention(self, query, key, value):
        score = tf.matmul(query, key, transpose_b=True)
        dim_key = tf.cast(tf.shape(key)[-1], tf.float32)
        scaled_score = score / tf.math.sqrt(dim_key)
        weights = tf.nn.softmax(scaled_score, axis=-1)
        out = tf.matmul(weights, value)
        return out, weights

    def separate_heads(self, x, batch_size):
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.projection_dim))
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, inputs):
        # x.shape = [batch_size, seq_len, num_variables]
        batch_size = tf.shape(inputs)[0]

        # Transpose inputs to (batch_size, num_variables, seq_len)
        inputs = tf.transpose(inputs, perm=[0, 2, 1])

        # Query
        query = self.query_dense(inputs)
        query = self.separate_heads(query, batch_size)

        # Key
        key = self.key_dense(inputs)
        key = self.separate_heads(key, batch_size)

        # Value
        value = self.value_dense(inputs)
        value = self.separate_heads(value, batch_size)

        # Attention
        query_attention, weights = self.attention(query, key, value)

        # Concatenate heads
        query_attention = tf.transpose(query_attention, perm=[0, 2, 1, 3])
        concat_attention = tf.reshape(query_attention, (batch_size, -1, self.embed_dim))

        # Output
        output = self.output_dense(concat_attention)

        return output

最后输出的张量形状为 [batch_size, seq_len, embed_dim]，其中 embed_dim 是指输入张量中每个时间步的特征维度。如果需要对时间步进行汇总，可以对第二维进行相应的汇聚操作。