1
# 传统方法：分段处理
2
segments = ["The cat sat on the mat", "It was a nice day"]
3
for segment in segments:
4
    # 每个段落独立处理，无跨段依赖
5
    output = model(segment)

1
class TransformerXLAttention(nn.Module):
2
    """Transformer-XL 的关键：跨段注意力"""
3

4
    def __init__(self, d_model, n_heads, segment_len):
5
        super().__init__()
6
        self.segment_len = segment_len
7
        self.h = {}  # 存储之前段的隐藏状态
8

9
    def forward(self, x, segment_id, layer_id):
10
        # x: [batch, seq_len, d_model]
11
        batch_size = x.size(0)
12

13
        # 获取当前段的序列长度
14
        curr_len = x.size(1)
15

16
        # 重建之前段的隐藏状态
17
        if layer_id not in self.h:
18
            # 第一个段，无历史状态
19
            prev_segments = []
20
        else:
21
            # 连接之前段的隐藏状态
22
            prev_segments = self.h[layer_id]
23

24
        # 关键创新：拼接历史状态
25
        if len(prev_segments) > 0:
26
            # 扩展上下文
27
            extended_x = torch.cat([prev_segments[-1], x], dim=1)
28
        else:
29
            extended_x = x
30

31
        # 标准注意力计算，但利用了更长的上下文
32
        attn_output = self.multi_head_attention(extended_x)
33

34
        # 更新存储的隐藏状态
35
        if segment_id not in self.h:
36
            self.h[segment_id] = []
37
        self.h[segment_id].append(x.detach())
38

39
        # 返回当前段的输出
40
        return attn_output[:, -curr_len:]

1
def relative_positional_encoding(query, key, rel_pos_embed):
2
    """
3
    相对位置编码的核心计算
4

5
    原始注意力: (S_j W_Q)^T (S_i W_K) = S_j^T W_Q^T W_K S_i
6

7
    引入相对位置后:
8
    A_{ij} = S_j^T W_Q^T W_K S_i + S_j^T W_Q^T W_R r_{i-j}
9

10
    其中 r_{i-j} 是相对位置嵌入
11
    """
12
    # 相对位置嵌入
13
    rel_pos = rel_pos_embed(torch.arange(-query_len, key_len))
14

15
    # 计算相对注意力分数
16
    rel_attn = torch.matmul(query, rel_pos.transpose(-2, -1))
17

18
    return rel_attn

1
class TransformerXLModel(nn.Module):
2
    """完整的 Transformer-XL 模型"""
3

4
    def __init__(self, vocab_size, d_model, n_layers, n_heads):
5
        super().__init__()
6
        self.embed = nn.Embedding(vocab_size, d_model)
7
        self.layers = nn.ModuleList([
8
            TransformerXLAttention(d_model, n_heads)
9
            for _ in range(n_layers)
10
        ])
11
        self.ln_f = nn.LayerNorm(d_model)
12
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
13

14
    def forward(self, input_ids, segment_id=None):
15
        x = self.embed(input_ids)
16

17
        for layer in self.layers:
18
            x = layer(x, segment_id, layer_id)
19

20
        x = self.ln_f(x)
21
        return self.lm_head(x)

1
# 训练时的关键技巧
2

3
# 1. 渐近段长度
4
segment_len = 64  # 开始
5
segment_len = 128  # 中期
6
segment_len = 256  # 最终
7

8
# 2. 记忆状态管理
9
def detach_memory(self):
10
    """分离记忆，防止梯度回传过远"""
11
    self.memory = [m.detach() for m in self.memory]
12

13
# 3. 动态评估
14
def eval_with_memory(self, x, memory):
15
    """使用记忆状态进行评估"""
16
    x = torch.cat([memory, x], dim=1)
17
    output = self.forward(x)
18
    new_memory = x[:, -self.mem_len:].detach()
19
    return output, new_memory