基于Transformer的时序算法Crossformer代码解析

Crossformer 是基于 Transformer 的时序建模算法，它通过创新的设计来处理时序数据，特别适用于长时间序列的预测任务。传统的 Transformer 模型通常在处理长时序数据时面临计算复杂度较高的问题，而 Crossformer 通过引入跨维度的注意力机制，有效地减少了计算负担，并提升了对时间序列中复杂模式的捕捉能力。该模型的设计不仅继承了 Transformer 的强大能力，还针对时序数据的特点进行了优化，使得它在多个时序预测任务中表现优异。

接下来，我将详细解析 Crossformer 的核心代码，具体包括DSW,TSA和HIERARCHICAL ENCODER-DECODER模块。

一、Dimension-Segment-Wise(DSW层)

1.1 问题背景:

• 传统 Transformer 嵌入方法将同一时间步上的所有维度嵌入为一个向量，关注跨时间依赖，但没有显式捕获跨维度依赖。

• 时间序列中的单点信息量较少，但其与邻近点形成的序列模式较为信息丰富。

• 注意力值表现出分段趋势，说明局部片段在时间序列建模中具有重要性。

1.2 提出方法:

• 每个嵌入向量不再表示单个时间步上的所有维度值，而是表示单维度上一段时间序列的片段。

• 将时间序列划分为长度为 τ 的片段，使用线性投影和位置嵌入得到嵌入向量。

1.3 代码实现

class DimeSegWise(nn.Module):
    def __init__(self, enc_in, in_seg_num, d_model, patch_len, stride, padding, dropout):
        super(DimeSegWise, self).__init__()
        # Patching: 设置片段长度和步幅
        self.patch_len = patch_len  # 每个时间序列片段的长度
        self.stride = stride  # 时间片段的滑动步长

        # 使用复制边界值的方式填充输入序列，以满足分段需求
        self.padding_patch_layer = nn.ReplicationPad1d((0, padding))

        # Backbone: 输入编码，投影到 d_model 维向量空间
        # 将每个片段映射到 d_model 维的特征空间
        self.value_embedding = nn.Linear(patch_len, d_model, bias=False)

        # 位置嵌入，用于为每个片段提供时间位置信息
        self.position_embedding = PositionalEmbedding(d_model)

        # 用于捕获跨维度依赖性的可学习位置嵌入
        # enc_pos_embedding 是一个可学习参数，形状为 [1, enc_in, in_seg_num, d_model]
        # enc_in: 输入序列的变量数，in_seg_num: 每个变量对应的片段数
        self.enc_pos_embedding = nn.Parameter(
            torch.randn(1, enc_in, in_seg_num, d_model)
        )

        # Dropout 层，用于随机失活部分特征，防止过拟合
        self.dropout = nn.Dropout(dropout)

        # 归一化层，用于对 d_model 维的特征进行规范化
        self.pre_norm = nn.LayerNorm(d_model)

    def forward(self, x):
        # 输入数据 x 的形状为 [B, V, L]，其中：
        # B: 批量大小（batch size），V: 变量数（维度数），L: 时间序列长度

        # 获取输入的变量数（维度数）
        n_vars = x.shape[1]

        # 对输入序列进行填充，确保可以按照给定的片段长度进行分割
        # 填充后形状仍为 [B, V, L+padding]
        x = self.padding_patch_layer(x)

        # 将时间序列分割成多个长度为 patch_len 的片段
        # unfold 操作将 x 的最后一维（时间维）按步长 stride 划分为小片段
        # 结果形状为 [B, V, P_NUM, P_LEN]，其中 P_NUM 是分割后的片段数
        x = x.unfold(dimension=-1, size=self.patch_len, step=self.stride)

        # 调整形状，将 [B, V, P_NUM, P_LEN] 转换为 [B * V, P_NUM, P_LEN]
        # 这样每个变量的片段会作为一个独立样本处理
        x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3]))

        # 对每个片段进行特征投影（value embedding），并加上位置嵌入
        # self.value_embedding(x) 将形状从 [B * V, P_NUM, P_LEN] 投影为 [B * V, P_NUM, d_model]
        # self.position_embedding(x) 提供位置编码，形状为 [1, P_NUM, d_model]
        x_enc = self.value_embedding(x) + self.position_embedding(x)

        # 恢复原始的变量维度，调整形状为 [B, V, P_NUM, d_model]
        # 其中 seg_num = P_NUM，d_model 是特征维度
        x_enc = rearrange(
            x_enc, '(b d) seg_num d_model -> b d seg_num d_model', d=n_vars
        )

        # 加上可学习的位置嵌入 enc_pos_embedding，用于捕获跨维度的依赖性
        # enc_pos_embedding 的形状为 [1, V, P_NUM, d_model]
        x_enc += self.enc_pos_embedding

        # 对特征进行归一化
        x_enc = self.pre_norm(x_enc)

        # 返回处理后的特征，形状为 [B, V, P_NUM, d_model]
        return x_enc

上述代码通过滑动窗口对时间序列进行分段。对每个时间片段进行投影到高维特征空间。加入时间和维度的嵌入信息，捕获跨时间和跨维度的依赖关系。

二、TWO-STAGE ATTENTION LAYER（TSA层）

两阶段注意力（Two-Stage Attention, TSA）层的主要作用是高效捕获二维向量数组中的 跨时间依赖性 和 跨维度依赖性，将时间序列和维度的复杂关系整合到统一的表示中。它为处理多元时间序列（MTS）的 Transformer 模型提供了一种针对性的改进，解决了直接应用自注意力导致计算复杂度过高的问题。

TSA 层分为两个阶段：跨时间阶段（Cross-Time Stage） 和 跨维度阶段（Cross-Dimension Stage）。

2.1 跨时间阶段

给定一个二维数组作为 TSA 层的输入，其中 L 和 D 分别表示段的数量和维度的数量。这里的 Z 可以是 DSW 嵌入的输出或更低 TSA 层的输出。为了方便描述，下文中用 表示时间步 i 的所有维度向量，用 Z:,d 表示维度 d 的所有时间步向量。

LayerNorm 表示层归一化，MLP 表示多层（本文为两层）前馈网络，MSA(Q, K, V) 表示多头自注意力层，其中 Q、K、V 分别作为查询、键和值。

直白理解：跨时间阶段是学习每个变量时间段之间的依赖关系。

2.2 跨维度阶段

跨维度阶段通过设计路由机制，高效捕获多元时间序列中维度间的依赖关系。路由机制在每个时间步中聚合所有维度的全局信息，并将其分配到各维度，实现全对全连接建模，显著降低计算复杂度（从O(D^2) 降至 O(DL)）。这一阶段解决了直接应用多头自注意力机制在高维场景下的计算瓶颈，同时确保不同维度间的信息交互，为高维多元时间序列数据的表示学习提供强大支持。

其中，（c为常数）是作为路由器的可学习向量数组。是来自所有维度的聚合信息。表示路由机制的输出。所有时间步共享相同的分别表示跳跃连接和 MLP 的输出。路由机制将复杂度从降低到O(DL)。

2.3 代码实现

class FullAttention(nn.Module):
    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
        super(FullAttention, self).__init__()
        # 初始化全注意力机制，支持掩码、缩放和注意力丢弃率
        self.scale = scale  # 控制缩放因子
        self.mask_flag = mask_flag  # 是否使用掩码
        self.output_attention = output_attention  # 是否输出注意力权重
        self.dropout = nn.Dropout(attention_dropout)  # 定义注意力输出的丢弃率

    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):
        # queries, keys, values: 输入的查询、键和值的张量
        # attn_mask: 可选的掩码，用于局部注意力或因果掩码
        B, L, H, E = queries.shape  # B: 批量大小, L: 序列长度, H: 注意力头数, E: 每头维度
        _, S, _, D = values.shape  # S: 键和值的序列长度, D: 每值维度
        scale = self.scale or 1. / sqrt(E)  # 设置缩放因子

        # 计算注意力分数，通过内积方法 (B, H, L, S)
        scores = torch.einsum("blhe,bshe->bhls", queries, keys)

        # 如果使用掩码，则应用掩码操作
        if self.mask_flag:
            if attn_mask is None:  # 如果未提供掩码，则创建因果掩码
                attn_mask = TriangularCausalMask(B, L, device=queries.device)
            scores.masked_fill_(attn_mask.mask, -np.inf)  # 掩码位置填充负无穷

        # 对分数应用 softmax，并加入 dropout 以提高鲁棒性
        A = self.dropout(torch.softmax(scale * scores, dim=-1))
        # 计算注意力加权输出 (B, L, H, D)
        V = torch.einsum("bhls,bshd->blhd", A, values)

        # 根据配置选择是否输出注意力权重
        if self.output_attention:
            return V.contiguous(), A
        else:
            return V.contiguous(), None

class AttentionLayer(nn.Module):
    def __init__(self, attention, d_model, n_heads, d_keys=None, d_values=None):
        super(AttentionLayer, self).__init__()
        # 初始化注意力层
        d_keys = d_keys or (d_model // n_heads)  # 查询和键的维度
        d_values = d_values or (d_model // n_heads)  # 值的维度

        self.inner_attention = attention  # 内部使用指定的注意力机制
        self.query_projection = nn.Linear(d_model, d_keys * n_heads)  # 查询的线性投影
        self.key_projection = nn.Linear(d_model, d_keys * n_heads)  # 键的线性投影
        self.value_projection = nn.Linear(d_model, d_values * n_heads)  # 值的线性投影
        self.out_projection = nn.Linear(d_values * n_heads, d_model)  # 输出线性变换
        self.n_heads = n_heads  # 注意力头数

    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):
        # 输入的查询、键和值分别通过线性投影并调整形状以适配多头注意力
        B, L, _ = queries.shape
        _, S, _ = keys.shape
        H = self.n_heads

        queries = self.query_projection(queries).view(B, L, H, -1)
        keys = self.key_projection(keys).view(B, S, H, -1)
        values = self.value_projection(values).view(B, S, H, -1)

        # 调用 FullAttention 计算注意力
        out, attn = self.inner_attention(
            queries,
            keys,
            values,
            attn_mask,
            tau=tau,
            delta=delta
        )
        out = out.view(B, L, -1)  # 合并多头输出

        # 返回经过线性投影的最终输出和注意力权重
        return self.out_projection(out), attn

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, configs, seg_num, factor, d_model, n_heads, d_ff=None, dropout=0.1):
        super(TwoStageAttentionLayer, self).__init__()
        # 初始化 TSA 层
        d_ff = d_ff or 4 * d_model  # 前馈网络的隐藏层大小

        # 定义跨时间和跨维度的注意力层，分别使用 FullAttention
        self.time_attention = AttentionLayer(FullAttention(False, configs.factor, attention_dropout=configs.dropout,
                                                           output_attention=configs.output_attention), d_model, n_heads)
        self.dim_sender = AttentionLayer(FullAttention(False, configs.factor, attention_dropout=configs.dropout,
                                                       output_attention=configs.output_attention), d_model, n_heads)
        self.dim_receiver = AttentionLayer(FullAttention(False, configs.factor, attention_dropout=configs.dropout,
                                                         output_attention=configs.output_attention), d_model, n_heads)
        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, attn_mask=None, tau=None, delta=None):
        # 输入 x 的形状: [batch, D, L, d_model]

        # Cross Time Stage
        # 将跨时间的输入调整形状，以适应时间维度上的自注意力
        batch = x.shape[0]
        time_in = rearrange(x, 'b ts_d seg_num d_model -> (b ts_d) seg_num d_model')

        # 对每个维度单独应用时间自注意力
        time_enc, attn = self.time_attention(time_in, time_in, time_in, attn_mask=None, tau=None, delta=None)
        dim_in = time_in + self.dropout(time_enc)  # 残差连接
        dim_in = self.norm1(dim_in)
        dim_in = dim_in + self.dropout(self.MLP1(dim_in))  # MLP 和残差连接
        dim_in = self.norm2(dim_in)

        # Cross Dimension Stage
        # 使用路由机制捕获维度之间的依赖关系
        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, attn = self.dim_sender(batch_router, dim_send, dim_send, attn_mask=None, tau=None, delta=None)
        # 然后分配信息到维度中
        dim_receive, attn = self.dim_receiver(dim_send, dim_buffer, dim_buffer, attn_mask=None, tau=None, delta=None)

        # 跨维度的残差连接和前馈网络
        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

三、HIERARCHICAL ENCODER-DECODER

3.1 Encoder

在编码器的每一层（除第一层外），时间域中每两个相邻的向量会被合并，以获得更粗略层级的表示。随后，应用 TSA 层来捕获此尺度上的依赖关系。此过程建模为 ：

其中，H 表示通过 DSW 嵌入获得的二维数组；表示第 l 层编码器的输出；表示用于段合并的可学习矩阵；表示拼接操作；表示第 l−1 层中每个维度的段数，如果不能被 2 整除，则对 进行填充以达到合适的长度；表示第 l 层段合并后的数组。假设编码器中有 N 层，用 表示编码器的 N+1 个输出。每层编码器的计算复杂度为 。

3.1.1 代码实现

class SegMerging(nn.Module):
    def __init__(self, d_model, win_size, norm_layer=nn.LayerNorm):
        """
        初始化SegMerging层，进行时间片段的合并操作。

        :param d_model: 特征维度大小
        :param win_size: 合并的窗口大小
        :param norm_layer: 归一化层，默认为LayerNorm
        """
        super().__init__()
        self.d_model = d_model  # 保存特征维度
        self.win_size = win_size  # 保存窗口大小
        self.linear_trans = nn.Linear(win_size * d_model, d_model)  # 合并后进行线性映射的层
        self.norm = norm_layer(win_size * d_model)  # 合并后归一化层

    def forward(self, x):
        """
        前向传播函数，执行时间片段合并操作。

        :param x: 输入张量，形状为 [b, v, seg_num, d_model]
            b: 批量大小
            v: 特征维度
            seg_num: 时间片段的数量
            d_model: 每个时间片段的特征维度
        :return: 输出张量，形状为 [b, v, seg_num // 2, d_model]
        """
        # 获取输入张量的形状：批量大小、时间片段数、每个时间片段的特征维度
        batch_size, ts_d, seg_num, d_model = x.shape

        # 如果时间片段数不是窗口大小的倍数，需要填充数据
        pad_num = seg_num % self.win_size
        if pad_num != 0:
            pad_num = self.win_size - pad_num
            # 将填充的数据沿时间片段维度拼接
            x = torch.cat((x, x[:, :, -pad_num:, :]), dim=-2)

        # 合并时间片段
        seg_to_merge = []
        for i in range(self.win_size):
            # 按照窗口大小取时间片段并合并
            seg_to_merge.append(x[:, :, i::self.win_size, :])

        # 将所有合并的时间片段拼接
        x = torch.cat(seg_to_merge, -1)

        # 对拼接后的数据进行归一化
        x = self.norm(x)

        # 对归一化后的数据进行线性映射，将其映射回原始维度
        x = self.linear_trans(x)

        return x


class scale_block(nn.Module):
    def __init__(self, configs, win_size, d_model, n_heads, d_ff, depth, dropout, 
                 seg_num=10, factor=10):
        """
        初始化scale_block层，其中包含了分段合并层和多个TSA层。

        :param configs: 配置参数
        :param win_size: 窗口大小
        :param d_model: 特征维度
        :param n_heads: 注意力头数
        :param d_ff: 前馈网络的隐藏层维度
        :param depth: TSA层的深度
        :param dropout: dropout比率
        :param seg_num: 每个时间片段的数量
        :param factor: 用于调节TSA层复杂度的因子
        """
        super(scale_block, self).__init__()

        # 如果窗口大小大于1，初始化分段合并层
        if win_size > 1:
            self.merge_layer = SegMerging(d_model, win_size, nn.LayerNorm)
        else:
            self.merge_layer = None

        # 初始化多个TSA层
        self.encode_layers = nn.ModuleList()

        for i in range(depth):
            # 每一层包含一个TwoStageAttentionLayer
            self.encode_layers.append(TwoStageAttentionLayer(configs, seg_num, factor, d_model, n_heads, 
                                                             d_ff, dropout))

    def forward(self, x, attn_mask=None, tau=None, delta=None):
        """
        前向传播函数，执行分段合并和TSA层的计算。

        :param x: 输入张量，形状为 [b, v, seg_num, d_model]
        :param attn_mask: 注意力掩码（默认为None）
        :param tau: 额外参数，用于调整TSA层（默认为None）
        :param delta: 额外参数，用于调整TSA层（默认为None）
        :return: 输出张量，形状为 [b, v, seg_num, d_model]
        """
        # 获取输入张量的时间维度
        _, ts_dim, _, _ = x.shape

        # 如果存在分段合并层，先执行分段合并
        if self.merge_layer is not None:
            x = self.merge_layer(x)

        # 按照TSA层的顺序进行前向传播
        for layer in self.encode_layers:
            x = layer(x)

        return x, None


class Encoder(nn.Module):
    def __init__(self, attn_layers):
        """
        初始化Encoder层，包含多个scale_block。

        :param attn_layers: 列表，包含多个scale_block
        """
        super(Encoder, self).__init__()
        self.encode_blocks = nn.ModuleList(attn_layers)

    def forward(self, x):
        """
        前向传播函数，依次通过所有的scale_block。

        :param x: 输入张量，形状为 [b, v, seg_num, d_model]
        :return: 输出张量的列表，每个元素表示在不同层的输出
        """
        encode_x = []
        encode_x.append(x)

        # 依次通过所有的scale_block
        for block in self.encode_blocks:
            x, attns = block(x)
            encode_x.append(x)

        return encode_x, None

SegMerging类：用于合并时间片段的操作。首先会判断输入的时间片段数是否能被窗口大小整除，如果不能，则填充时间片段数。接着将每个窗口内的片段进行合并，最终经过线性转换和归一化，得到合并后的输出。

scale_block类：用于构建包含多个TSA层的块。在每个块中，如果窗口大小大于1，首先会应用SegMerging进行时间片段合并。接着，会按顺序执行多个TwoStageAttentionLayer层，这些层用于捕捉时间序列的跨时间和跨维度的依赖关系。

Encoder类：包含多个scale_block。它将输入数据逐层传递通过每个scale_block，并在每一层输出时保存输出结果。最终，返回所有层的输出。

3.2 Decoder

获取由编码器输出的 N+1 个特征数组后，解码器中使用 N+1 层（索引为 0,1,…,N）进行预测。第 l 层以第 l 个编码数组为输入，然后输出第 l 层的解码二维数组。该过程总结为:

其中，表示解码器的可学习位置嵌入。是 TSA 层的输出。MSA 层将 作为查询，作为键和值，用于建立编码器和解码器之间的连接。MSA 的输出记为 。分别表示跳跃连接和 MLP 的输出。用 来表示解码器的输出。每个解码器层的计算复杂度为：。

线性投影应用于每一层的输出，以生成该层的预测。所有层的预测结果相加以得到最终的预测（对于 l=0,…,N）：

其中 是一个可学习的矩阵，用于将向量投影到一个时间序列段。表示预测中第 i 个段在维度 d 上的值。层 l 的所有段被重新排列，得到该层的预测 。所有层的预测结果相加以得到最终的预测 。

3.2 代码实现

class AttentionLayer(nn.Module):
    """
    标准的 Attention 层，应用多头注意力机制到输入的查询（queries）、键（keys）和值（values）。
    包含对查询、键和值的线性变换，并进行注意力计算。

    参数：
    - attention: 使用的注意力机制（例如 FullAttention 等）
    - d_model: 输入特征空间的维度
    - n_heads: 注意力头的数量
    - d_keys: 键向量的维度，默认值为 d_model // n_heads
    - d_values: 值向量的维度，默认值为 d_model // n_heads
    """
    def __init__(self, attention, d_model, n_heads, d_keys=None, d_values=None):
        super(AttentionLayer, self).__init__()

        # 默认情况下，d_keys 和 d_values 被设置为 d_model // n_heads
        d_keys = d_keys or (d_model // n_heads)
        d_values = d_values or (d_model // n_heads)

        self.inner_attention = attention  # 内部的注意力机制（例如 FullAttention）
        # 将输入的 d_model 映射到键和值的维度
        self.query_projection = nn.Linear(d_model, d_keys * n_heads)
        self.key_projection = nn.Linear(d_model, d_keys * n_heads)
        self.value_projection = nn.Linear(d_model, d_values * n_heads)
        # 最终的线性变换，将输出映射回 d_model 空间
        self.out_projection = nn.Linear(d_values * n_heads, d_model)
        self.n_heads = n_heads  # 注意力头的数量

    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):
        """
        Attention 层的前向传播，执行多头注意力计算。

        参数：
        - queries: 输入查询，形状为 [batch_size, sequence_length, d_model]
        - keys: 输入键，形状为 [batch_size, sequence_length, d_model]
        - values: 输入值，形状为 [batch_size, sequence_length, d_model]
        - attn_mask: 注意力掩码，用于指定哪些位置在计算中应被忽略
        - tau, delta: 可选的额外参数，具体取决于注意力机制的实现

        返回：
        - out: 通过注意力机制的输出，形状为 [batch_size, sequence_length, d_model]
        - attn: 注意力权重，用于分析注意力分布
        """
        B, L, _ = queries.shape  # 获取批大小（batch size）和查询序列长度
        _, S, _ = keys.shape  # 获取键的序列长度
        H = self.n_heads  # 获取注意力头的数量

        # 线性变换：将查询、键和值映射到多个注意力头的维度
        queries = self.query_projection(queries).view(B, L, H, -1)
        keys = self.key_projection(keys).view(B, S, H, -1)
        values = self.value_projection(values).view(B, S, H, -1)

        # 通过内部注意力机制计算注意力输出和注意力权重
        out, attn = self.inner_attention(
            queries,
            keys,
            values,
            attn_mask,
            tau=tau,
            delta=delta
        )

        # 将输出的形状转换回 [batch_size, sequence_length, -1]
        out = out.view(B, L, -1)

        return self.out_projection(out), attn  # 最后通过输出层进行映射并返回

class TwoStageAttentionLayer(nn.Module):
    """
    两阶段注意力层（TSA）：
    该层包括两个阶段的注意力计算：跨时间阶段（Cross Time Stage）和跨维度阶段（Cross Dimension Stage）。
    输入和输出的形状为 [batch_size, Data_dim(D), Seg_num(L), d_model]

    参数：
    - configs: 配置信息，包含各类超参数
    - seg_num: 分段数
    - factor: 用于控制注意力机制的因子
    - d_model: 输入的特征维度
    - n_heads: 注意力头的数量
    - d_ff: 前馈网络的维度，默认为 4 * d_model
    - dropout: Dropout 概率，默认为 0.1
    """
    def __init__(self, configs, 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  # 默认的前馈网络维度为 d_model 的四倍
        # 初始化两个阶段的注意力机制
        self.time_attention = AttentionLayer(
            FullAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=configs.output_attention),
            d_model, n_heads
        )
        self.dim_sender = AttentionLayer(
            FullAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=configs.output_attention),
            d_model, n_heads
        )
        self.dim_receiver = AttentionLayer(
            FullAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=configs.output_attention),
            d_model, n_heads
        )
        self.router = nn.Parameter(torch.randn(seg_num, factor, d_model))  # 学习的路由向量

        self.dropout = nn.Dropout(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, attn_mask=None, tau=None, delta=None):
        """
        前向传播，分为跨时间阶段（Cross Time Stage）和跨维度阶段（Cross Dimension Stage）两步。

        参数：
        - x: 输入数据，形状为 [batch_size, seg_num, d_model]
        - attn_mask: 可选的注意力掩码
        - tau, delta: 可选的附加参数，用于注意力机制

        返回：
        - final_out: 最终输出，形状为 [batch_size, seg_num, d_model]
        """
        batch = x.shape[0]

        # 跨时间阶段：直接对每个维度应用 MSA（多头自注意力）
        time_in = rearrange(x, 'b ts_d seg_num d_model -> (b ts_d) seg_num d_model')
        time_enc, attn = self.time_attention(time_in, time_in, time_in, attn_mask=None, tau=None, delta=None)
        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)
        batch_router = repeat(self.router, 'seg_num factor d_model -> (repeat seg_num) factor d_model', repeat=batch)
        dim_buffer, attn = self.dim_sender(batch_router, dim_send, dim_send, attn_mask=None, tau=None, delta=None)
        dim_receive, attn = self.dim_receiver(dim_send, dim_buffer, dim_buffer, attn_mask=None, tau=None, delta=None)
        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)

        # 最终输出，重新排列为 [batch_size, seg_num, d_model]
        final_out = rearrange(dim_enc, '(b seg_num) ts_d d_model -> b ts_d seg_num d_model', b=batch)

        return final_out


class DecoderLayer(nn.Module):
    """
    解码器层，包括自注意力（self-attention）和跨注意力（cross-attention）机制。

    参数：
    - self_attention: 自注意力层
    - cross_attention: 跨注意力层
    - seg_len: 每个时间段的长度
    - d_model: 输入特征维度
    - d_ff: 前馈网络的维度
    - dropout: Dropout 概率，默认为 0.1
    """
    def __init__(self, self_attention, cross_attention, seg_len, d_model, d_ff=None, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.self_attention = self_attention
        self.cross_attention = cross_attention
        self.seg_len = seg_len
        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, context, attn_mask=None, tau=None, delta=None):
        """
        前向传播，包括自注意力和跨注意力机制。

        参数：
        - x: 输入数据
        - context: 额外上下文数据，用于跨注意力计算
        - attn_mask: 可选的注意力掩码
        - tau, delta: 可选的附加参数

        返回：
        - final_out: 最终输出
        """
        # 自注意力层
        self_attn_out, attn = self.self_attention(x, x, x, attn_mask=attn_mask, tau=tau, delta=delta)
        self_attn_out = self.norm1(self_attn_out)
        self_attn_out = self_attn_out + self.dropout(self.MLP1(self_attn_out))
        self_attn_out = self.norm2(self_attn_out)

        # 跨注意力层
        cross_attn_out, attn = self.cross_attention(self_attn_out, context, context, attn_mask=attn_mask, tau=tau, delta=delta)
        cross_attn_out = self.norm3(cross_attn_out)
        cross_attn_out = cross_attn_out + self.dropout(self.MLP2(cross_attn_out))
        cross_attn_out = self.norm4(cross_attn_out)

        return cross_attn_out

AttentionLayer 是一个多头注意力机制的实现，它通过将输入映射到多个子空间并并行计算多个注意力头，来捕捉输入序列中不同部分之间的依赖关系。最终，通过合并这些注意力头的输出并进行线性变换，得到最终的注意力结果。

TwoStageAttentionLayer 是一个两阶段注意力机制，首先通过跨时间阶段进行自注意力计算来捕捉时间序列的依赖关系，然后通过跨维度阶段使用路由向量来聚合不同维度的信息，从而增强模型的多维信息处理能力。

DecoderLayer 是 Transformer 解码器中的基本组件，它包含自注意力层和跨注意力层。自注意力层处理输入序列的内部依赖关系，跨注意力层则结合来自编码器的上下文信息，用以生成最终的输出序列。该层还包括一个前馈神经网络进一步处理注意力结果。

四、完整代码

完整代码：https://github.com/thuml/Time-Series-Library/blob/main/models/Crossformer.py

Crossformer，一种基于 Transformer 的模型，利用跨维度依赖性进行多变量时间序列（MTS）预测。具体而言，Dimension-Segment-Wise (DSW) 嵌入将输入数据嵌入为一个 2D 向量数组，以保留时间和维度的相关信息。设计了 Two-Stage-Attention (TSA) 层，以捕捉嵌入数组的跨时间和跨维度依赖性。通过使用 DSW 嵌入和 TSA 层，设计了一个层次化编码器-解码器（HED），以利用不同尺度的信息。在六个真实世界数据集上的实验结果表明，Crossformer 优于以往的最先进方法。

---

五、如何系统学习掌握AI大模型？

AI大模型作为人工智能领域的重要技术突破，正成为推动各行各业创新和转型的关键力量。抓住AI大模型的风口，掌握AI大模型的知识和技能将变得越来越重要。

学习AI大模型是一个系统的过程，需要从基础开始，逐步深入到更高级的技术。

这里给大家精心整理了一份全面的AI大模型学习资源，包括：AI大模型全套学习路线图（从入门到实战）、精品AI大模型学习书籍手册、视频教程、实战学习、面试题等，资料免费分享！

1. 成长路线图&学习规划

要学习一门新的技术，作为新手一定要先学习成长路线图，方向不对，努力白费。

这里，我们为新手和想要进一步提升的专业人士准备了一份详细的学习成长路线图和规划。可以说是最科学最系统的学习成长路线。

2. 大模型经典PDF书籍

书籍和学习文档资料是学习大模型过程中必不可少的，我们精选了一系列深入探讨大模型技术的书籍和学习文档，它们由领域内的顶尖专家撰写，内容全面、深入、详尽，为你学习大模型提供坚实的理论基础。（书籍含电子版PDF）

3. 大模型视频教程

对于很多自学或者没有基础的同学来说，书籍这些纯文字类的学习教材会觉得比较晦涩难以理解，因此，我们提供了丰富的大模型视频教程，以动态、形象的方式展示技术概念，帮助你更快、更轻松地掌握核心知识。

4. 2024行业报告

行业分析主要包括对不同行业的现状、趋势、问题、机会等进行系统地调研和评估，以了解哪些行业更适合引入大模型的技术和应用，以及在哪些方面可以发挥大模型的优势。

5. 大模型项目实战

学以致用 ，当你的理论知识积累到一定程度，就需要通过项目实战，在实际操作中检验和巩固你所学到的知识，同时为你找工作和职业发展打下坚实的基础。

6. 大模型面试题

面试不仅是技术的较量，更需要充分的准备。

在你已经掌握了大模型技术之后，就需要开始准备面试，我们将提供精心整理的大模型面试题库，涵盖当前面试中可能遇到的各种技术问题，让你在面试中游刃有余。

全套的AI大模型学习资源已经整理打包，有需要的小伙伴可以微信扫描下方优快云官方认证二维码，免费领取【保证100%免费】