【ICLR 2023】时序精选论文01｜CrossFormer: 实现多维变量间的羁绊（代码解读附源码）

提出此方案主要解决的问题点：

作者认为先前的基于Transformer的模型在捕获长期时间依赖性上可谓是下足了功夫，还提出各种Attention变体来降低复杂度。然而，这些方法都忽略了对不同变量的“跨维度依赖性”的捕获，也就是多变量时序中不同变量序列的相关性。作者认为对这种相关性的捕获是重要的，尤其是对于多变量时序预测任务来说。

大框架上仍然使用了Encoder-Decoder结构

Encoder

在代码中，Encoder部分先以一个scale_block进行封装，内部包含attention层，block块如下：

class Encoder(nn.Module):
    '''
    The Encoder of Crossformer.
    '''
    def __init__(self, e_blocks, win_size, d_model, n_heads, d_ff, block_depth, dropout,
                in_seg_num = 10, factor=10):
        super(Encoder, self).__init__()
        self.encode_blocks = nn.ModuleList()

        self.encode_blocks.append(scale_block(1, d_model, n_heads, d_ff, block_depth, dropout,\
                                            in_seg_num, factor))
        for i in range(1, e_blocks):
            self.encode_blocks.append(scale_block(win_size, d_model, n_heads, d_ff, block_depth, dropout,\
                                            ceil(in_seg_num/win_size**i), factor))

    def forward(self, x):
        encode_x = []
        encode_x.append(x)
        
        for block in self.encode_blocks:
            x = block(x)
            encode_x.append(x)

        return encode_x

block是由TwoStageAttention构成

class TwoStageAttentionLayer(nn.Module):
    '''
    The Two Stage Attention (TSA) Layer
    input/output shape: [batch_size, Data_dim(D), Seg_num(L), d_model]
    '''
    def __init__(self, seg_num, factor, d_model, n_heads, d_ff = None, dropout=0.1):
        super(TwoStageAttentionLayer, self).__init__()
        d_ff = d_ff or 4*d_model
        self.time_attention = AttentionLayer(d_model, n_heads, dropout = dropout)
        self.dim_sender = AttentionLayer(d_model, n_heads, dropout = dropout)
        self.dim_receiver = AttentionLayer(d_model, n_heads, dropout = dropout)
        self.router = nn.Parameter(torch.randn(seg_num, factor, d_model))
        
        self.dropout = nn.Dropout(dropout)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.norm4 = nn.LayerNorm(d_model)

        self.MLP1 = nn.Sequential(nn.Linear(d_model, d_ff),
                                nn.GELU(),
                                nn.Linear(d_ff, d_model))
        self.MLP2 = nn.Sequential(nn.Linear(d_model, d_ff),
                                nn.GELU(),
                                nn.Linear(d_ff, d_model))

    def forward(self, x):

        # Cross Time Stage: Directly apply MSA to each dimension
        batch = x.shape[0]
        time_in = rearrange(x, 'b ts_d seg_num d_model -> (b ts_d) seg_num d_model')
        time_enc = self.time_attention(
            time_in, time_in, time_in
        )
        dim_in = time_in + self.dropout(time_enc)
        dim_in = self.norm1(dim_in)
        dim_in = dim_in + self.dropout(self.MLP1(dim_in))
        dim_in = self.norm2(dim_in)

        # Cross Dimension Stage: use a small set of learnable vectors
        # to aggregate and distribute messages to build the D-to-D connection
        dim_send = rearrange(dim_in, '(b ts_d) seg_num d_model -> (b seg_num) ts_d d_model', b = batch)
        batch_router = repeat(self.router, 'seg_num factor d_model -> (repeat seg_num) factor d_model', repeat = batch)
        dim_buffer = self.dim_sender(batch_router, dim_send, dim_send)
        dim_receive = self.dim_receiver(dim_send, dim_buffer, dim_buffer)
        dim_enc = dim_send + self.dropout(dim_receive)
        dim_enc = self.norm3(dim_enc)
        dim_enc = dim_enc + self.dropout(self.MLP2(dim_enc))
        dim_enc = self.norm4(dim_enc)

        final_out = rearrange(dim_enc, '(b seg_num) ts_d d_model -> b ts_d seg_num d_model', b = batch)

        return final_out

此部分是关键点，CrossFormer中所使用的注意力构成为两阶段注意力，其中分别为时间维度上的注意力机制与特征维度上的注意力机制。与以往transformer不同的处理方式是，这种两阶段的注意力机制从两个不同的维度对输入进行计算权重。

起始数据输入维度为 (b,ts_d,input), 注意，这里作者在使用embedding的时候，先将单个样本进行了切分，固定长度下分割为了数等份，设定小片段的长度为seg_len，将每个样本分成小片段，形状为(b,ts_d,seg_num,seg_len),经过两层embedding之后，得到形状为(b,ts_d,seg_num,d_model)。这里作者认为这种方式可以让模型更好学习片段内的信息。

如图，这种处理方式在论文中被称为DSW，不过这种方式为什么就有效，目前还没有想透彻。

经过映射后，就开始进入Encoder计算，作者这里先将数据形状更换为如下，将样本序列这个维度放出来，进行注意力机制计算，然后进入时间注意力模块计算，其实这里的注意力机制都是使用的自注意力机制。

time_in = rearrange(x, 'b ts_d seg_num d_model -> (b ts_d) seg_num d_model')

time-attention输出的数据形状仍然和time_in相同，但已经经过了权重映射。

然后time_in再经过一个残差块的计算得到残差输出。而后作者在这里又加入了一个MLP1的模块（这里可能也是为了结构更深吧）

dim_in = time_in + self.dropout(time_enc)
dim_in = self.norm1(dim_in)
dim_in = dim_in + self.dropout(self.MLP1(dim_in))
dim_in = self.norm2(dim_in)

而后，进入二阶段注意力机制的第二阶段注意力，维度注意力机制

dim_send = rearrange(dim_in, '(b ts_d) seg_num d_model -> (b seg_num) ts_d d_model', b = batch)

这里的话，将dim_in的维度重新排列，并将维度这一列放出来，进行注意力计算。

batch_router = repeat(self.router, 'seg_num factor d_model -> (repeat seg_num) factor d_model', repeat = batch)
dim_buffer = self.dim_sender(batch_router, dim_send, dim_send)
dim_receive = self.dim_receiver(dim_send, dim_buffer, dim_buffer)

这里的话，有一个很关键的点，这里设置了一个中间变量，类似于VAE的特征映射操作情况，把复杂不可知的数据分布通过深度学习网络映射到非常浓缩的低维特征然后去重构原始输入。放在这就是把多通道嵌入的信息以不同方式浓缩到几个不同factor里，每个factor再决定把多少自己的信息转发给各个通道的嵌入数据。

VAE的结构图

方式很新颖，不过因为factor数量确实要控制得非常少的话，其效果大概率不能和使用完整的自注意力层来比，这中间可能出现较多的信息损失。

随后，再经过残差结构以及MLP2的计算，最终得到网络输出final_out

final_out = rearrange(dim_enc, '(b seg_num) ts_d d_model -> b ts_d seg_num d_model', b = batch)

整体的编码器结构：

Encoder里面包含e_blocks个scale_block，

一个scale_block里面包含depth个TwoStageAttentionLayer。

SegMerging结构

注意，这里有个SegMerging结构，在第二个scale_block开始使用，数据输入形状为[B, ts_d, L, d_model]。在第二个block计算时，则采用了这种片段融合结构，这里主要的一个参数为win_size，意义为将几个片段融合为一个，得到的数据输出为 [B, ts_d, seg_num/win_size, win_size*d_model]，再经过一个映射层，将其数据形状变换为[B, ts_d, seg_num/win_size, d_model]。

加入分段融合结构的原因，可能并不完全是出于减少计算复杂度的简单考虑，更多的是为了通过结构化的方式增强模型对数据的把握，包括更好的上下文理解和信息的有效蒸馏。这种设计允许模型在减少细节的同时，更有效地捕捉和利用关键信息，从而在多个层面上优化模型的性能和效率。

至此，整个编码器的结构完成，最终返回一个列表encode_x，包含每一个block的编码输出。

def forward(self, x):
    encode_x = []
    encode_x.append(x)
    
    for block in self.encode_blocks:
        x = block(x)
        encode_x.append(x)

    return encode_x

Decoder

数据经过解码器之后，解码器需要两个输入：编码器输出、构建的解码器输入

dec_in = repeat(self.dec_pos_embedding, 'b ts_d l d -> (repeat b) ts_d l d', repeat = batch_size)
predict_y = self.decoder(dec_in, enc_out)

解码器结构主要d_layers层解码器层DecoderLayer构成，每个DecoderLayer的结构是相同的

batch = x.shape[0]
x = self.self_attention(x)
x = rearrange(x, 'b ts_d out_seg_num d_model -> (b ts_d) out_seg_num d_model')

cross = rearrange(cross, 'b ts_d in_seg_num d_model -> (b ts_d) in_seg_num d_model')
tmp = self.cross_attention(
    x, cross, cross,
)
x = x + self.dropout(tmp)
y = x = self.norm1(x)
y = self.MLP1(y)
dec_output = self.norm2(x+y)

dec_output = rearrange(dec_output, '(b ts_d) seg_dec_num d_model -> b ts_d seg_dec_num d_model', b = batch)
layer_predict = self.linear_pred(dec_output)
layer_predict = rearrange(layer_predict, 'b out_d seg_num seg_len -> b (out_d seg_num) seg_len')

此处x为解码器输入dec_in，cross为编码器输出enc_out（列表存储了每个block的输出），x进入self_attention进行注意力（二阶段注意力）计算，然后再将x（已进行注意力加权）与cross进行注意力（全注意力）计算。再进行一个残差结构后，再经过一个norm层后进入MLP层计算，再经过一个残差结构得到输出dec_output。最后经过线性层得到输出b, ts_d, seg_dec_num, seg_len

def forward(self, x, cross):
    final_predict = None
    i = 0

    ts_d = x.shape[1]
    for layer in self.decode_layers:
        cross_enc = cross[i]
        x, layer_predict = layer(x,  cross_enc)
        if final_predict is None:
            final_predict = layer_predict
        else:
            final_predict = final_predict + layer_predict
        i += 1

可以看到，在解码器的forward中，由d_layers个解码器层构成，在每次的输入中，解码器输出的x，即dec_output仍然作为下一层的DecoderLayer的解码器输入。而编码器的输出在放入DecoderLayer中时，之前提到过enc_out是一个列表，每次都将对应的enc_out放入DecoderLayer中。关于输出，每一层的DecoderLayer都会有两个输出，一个是为了继续输入到下一层的DecoderLayer中，另外一个是我们需要的输出，且最终的结果为每一层的layer_predict相叠加得到的。

解码器中输入多层编码器输出的结构的优势：

自己的理解：

1、在深层网络中，仅使用最后一层的输出可能会导致梯度消失或爆炸问题，尤其是在非常深的网络中。将每一层的输出都用作解码器的输入可以帮助改善梯度的传播，使梯度更直接地流向网络的早期层次，从而提高训练过程的稳定性和效率。且在模型输出时，也是将每一层的解码器输出进行相加，这有益于梯度反向计算时的梯度传递，也在一定程度上增强了训练稳定性。

2、增强特征捕捉能力，每一层的编码器可能会捕捉到不同层次和不同类型的特征。在早期的层次中，编码器可能更多地关注于输入数据的局部特征或细节信息，而在更深层次中，则能捕捉到更抽象的高级特征。将所有这些层的输出提供给解码器可以让解码器访问从局部到全局的各种特征，从而提高模型的表达能力和最终的性能。

注意，这里的decoder层数设置上比encoder层数上要多1，从Encoder部分，假设我们设置了N层Encoder，那么算上初始嵌入、每层输出，我们就会得到N+1组嵌入表示，其中，初始嵌入和第1层的Encoder输出是相同维度。Crossformer对应着设置了N+1层解码器.

最后得到网络输出，解码器结构至此结束。

源码下载：基于CrossFormer的时间序列预测代码