目录
10.1.1 RNN的BPTT算法与梯度消失/爆炸的谱分析，LSTM的门控机制与GRU的简化门控结构
10.1.2 Seq2Seq架构与教师强制的暴露偏差问题，计划采样的退火策略
10.1.3 双向RNN与多层堆叠的并行化实现，CuDNN内核优化原理
10.2.1 Bahdanau注意力与Luong注意力的分数函数对比，键-查询-值矩阵乘法视角
10.2.2 多头自注意力的多头分割与输出拼接，缩放点积的Softmax温度缩放与注意力权重稀疏性分析
10.2.3 位置编码：正弦/余弦绝对位置编码的波长设计，可学习位置编码与RoPE旋转位置编码的复数域实现
10.2.4 Transformer编码器-解码器架构：残差连接的Pre-LN与Post-LN布局对比，FFN的GELU激活与中间维度扩展比例
10.2.5 高效注意力变体：Sparse Transformer的稀疏模式、Linear Attention的核技巧、Flash Attention的IO感知分块与Softmax在线统计量计算
10.1 循环神经网络与长程依赖
10.1.1 RNN的BPTT算法与梯度消失/爆炸的谱分析，LSTM的门控机制与GRU的简化门控结构
循环神经网络通过时间反向传播算法进行训练，该算法将计算图在时间维度上展开，应用链式法则计算参数的梯度。然而，当序列长度增加时，梯度需要通过多个时间步反向传播，导致梯度矩阵的奇异值连乘。当最大奇异值大于一时，梯度幅度随层深指数增长，引发梯度爆炸；当最大奇异值小于一时，梯度幅度指数衰减，导致梯度消失。这种谱特性的不稳定性使得标准RNN难以捕捉超过十个时间步的长程依赖关系。
长短期记忆网络通过引入细胞状态和三个门控机制解决此问题。遗忘门通过sigmoid激活决定历史信息的保留比例，输入门控制当前输入向细胞状态的写入程度，输出门调节细胞状态向隐藏状态的暴露水平。这种门控结构通过加法更新替代乘法更新，创建了一条梯度流动的高速公路，使得误差信号可以在 hundreds of time steps 中保持恒定幅度，从而有效缓解梯度消失。
门控循环单元对LSTM进行了结构性简化，将遗忘门与输入门耦合为单一的更新门，同时合并细胞状态与隐藏状态。这种设计将参数数量减少约百分之二十五，在保持相近建模能力的同时提升了计算效率。更新门决定历史状态与候选状态的混合比例，重置门控制历史信息在计算候选状态时的参与度，形成了更为紧凑的门控表达式。
脚本：rnn_gating_mechanisms.py
Python
复制
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脚本内容：RNN、LSTM、GRU底层机制实现与梯度流分析
使用方式：直接运行 python rnn_gating_mechanisms.py
依赖包：numpy, matplotlib, torch
功能说明：
1. 手动实现Vanilla RNN、LSTM、GRU的前向与反向传播
2. 展示BPTT中的梯度消失与爆炸现象（谱分析视角）
3. 对比门控机制对梯度流的影响
4. 可视化长程依赖学习中的梯度幅度衰减曲线
"""
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from typing import Tuple, List, Dict
import copy
class ManualRNN:
"""手动实现的Vanilla RNN，用于展示BPTT和梯度动态"""
def __init__(self, input_dim: int, hidden_dim: int, output_dim: int,
init_scale: float = 0.01):
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.output_dim = output_dim
# Xavier初始化
self.Wxh = np.random.randn(hidden_dim, input_dim) * init_scale
self.Whh = np.random.randn(hidden_dim, hidden_dim) * init_scale
self.Why = np.random.randn(output_dim, hidden_dim) * init_scale
self.bh = np.zeros((hidden_dim, 1))
self.by = np.zeros((output_dim, 1))
# 存储中间状态用于BPTT
self.hs = {}
self.xs = {}
self.ys = {}
self.ps = {}
def forward(self, inputs: List[np.ndarray]) -> List[np.ndarray]:
"""前向传播，返回输出概率分布序列"""
h = np.zeros((self.hidden_dim, 1))
self.hs[-1] = h
outputs = []
for t, x in enumerate(inputs):
self.xs[t] = x.reshape(-1, 1)
# 隐藏状态更新：tanh(W_xh @ x + W_hh @ h + b_h)
h = np.tanh(self.Wxh @ self.xs[t] + self.Whh @ h + self.bh)
self.hs[t] = h
# 输出层
y = self.Why @ h + self.by
self.ys[t] = y
# softmax概率
p = np.exp(y) / np.sum(np.exp(y))
self.ps[t] = p
outputs.append(p)
return outputs
def backward(self, targets: List[int], learning_rate: float = 0.1) -> Dict[str, np.ndarray]:
"""BPTT算法实现，返回梯度范数用于分析"""
dWxh = np.zeros_like(self.Wxh)
dWhh = np.zeros_like(self.Whh)
dWhy = np.zeros_like(self.Why)
dbh = np.zeros_like(self.bh)
dby = np.zeros_like(self.by)
dh_next = np.zeros((self.hidden_dim, 1))
grad_norms = [] # 记录每时间步的梯度范数
for t in reversed(range(len(targets))):
# 输出层梯度
dy = copy.deepcopy(self.ps[t])
dy[targets[t]] -= 1
dWhy += dy @ self.hs[t].T
dby += dy
# 反向传播到隐藏层
dh = self.Why.T @ dy + dh_next
# tanh梯度
dh_raw = dh * (1 - self.hs[t] ** 2)
dbh += dh_raw
dWxh += dh_raw @ self.xs[t].T
dWhh += dh_raw @ self.hs[t-1].T
# 下一时间步的隐藏层梯度
dh_next = self.Whh.T @ dh_raw
# 记录梯度范数（谱分析视角）
grad_norm = np.linalg.norm(dh_next, 2)
grad_norms.append(grad_norm)
# 梯度裁剪（防止爆炸）
for grad in [dWxh, dWhh, dWhy]:
np.clip(grad, -5, 5, out=grad)
# 参数更新
self.Wxh -= learning_rate * dWxh
self.Whh -= learning_rate * dWhh
self.Why -= learning_rate * dWhy
self.bh -= learning_rate * dbh
self.by -= learning_rate * dby
return {'grad_norms': list(reversed(grad_norms))}
class ManualLSTM:
"""手动实现的LSTM，展示门控机制对梯度流的保护"""
def __init__(self, input_dim: int, hidden_dim: int, init_scale: float = 0.01):
self.hidden_dim = hidden_dim
# 门控参数：遗忘门、输入门、输出门、候选细胞状态
self.Wf = np.random.randn(hidden_dim, input_dim + hidden_dim) * init_scale
self.Wi = np.random.randn(hidden_dim, input_dim + hidden_dim) * init_scale
self.Wo = np.random.randn(hidden_dim, input_dim + hidden_dim) * init_scale
self.Wc = np.random.randn(hidden_dim, input_dim + hidden_dim) * init_scale
self.bf = np.zeros((hidden_dim, 1))
self.bi = np.zeros((hidden_dim, 1))
self.bo = np.zeros((hidden_dim, 1))
self.bc = np.zeros((hidden_dim, 1))
# 存储中间状态
self.states = {}
def _sigmoid(self, x):
return 1 / (1 + np.exp(-np.clip(x, -500, 500)))
def forward(self, inputs: List[np.ndarray]) -> List[np.ndarray]:
"""LSTM前向传播，返回隐藏状态序列"""
h_prev = np.zeros((self.hidden_dim, 1))
c_prev = np.zeros((self.hidden_dim, 1))
self.states = {'hs': {-1: h_prev}, 'cs': {-1: c_prev},
'gates': {}, 'candidates': {}}
outputs = []
for t, x in enumerate(inputs):
# 拼接输入和前一隐藏状态
concat = np.vstack((h_prev, x.reshape(-1, 1)))
# 门控计算
f_t = self._sigmoid(self.Wf @ concat + self.bf) # 遗忘门
i_t = self._sigmoid(self.Wi @ concat + self.bi) # 输入门
o_t = self._sigmoid(self.Wo @ concat + self.bo) # 输出门
c_tilde = np.tanh(self.Wc @ concat + self.bc) # 候选细胞状态
# 细胞状态更新（关键的高速公路连接）
c_t = f_t * c_prev + i_t * c_tilde
# 隐藏状态更新
h_t = o_t * np.tanh(c_t)
self.states['gates'][t] = {'f': f_t, 'i': i_t, 'o': o_t, 'c_tilde': c_tilde}
self.states['cs'][t] = c_t
self.states['hs'][t] = h_t
outputs.append(h_t)
h_prev, c_prev = h_t, c_t
return outputs
def analyze_gradient_flow(self, seq_length: int) -> List[float]:
"""
分析梯度流经细胞状态的稳定性
返回各时间步的梯度范数
"""
grad_norms = []
# 模拟从最后一个时间步反向传播梯度
dc_next = np.random.randn(self.hidden_dim, 1) * 0.1
for t in reversed(range(seq_length)):
gates = self.states['gates'][t]
# 通过细胞状态的梯度流（加法路径保持梯度稳定）
dc_t = dc_next * gates['f'] # 遗忘门调制
grad_norm = np.linalg.norm(dc_t, 2)
grad_norms.append(grad_norm)
dc_next = dc_t
return list(reversed(grad_norms))
class ManualGRU:
"""手动实现的GRU，展示简化门控结构"""
def __init__(self, input_dim: int, hidden_dim: int, init_scale: float = 0.01):
self.hidden_dim = hidden_dim
# GRU只有两个门：重置门和更新门，合并了候选状态计算
self.Wr = np.random.randn(hidden_dim, input_dim + hidden_dim) * init_scale
self.Wz = np.random.randn(hidden_dim, input_dim + hidden_dim) * init_scale
self.Wh = np.random.randn(hidden_dim, input_dim + hidden_dim) * init_scale
self.br = np.zeros((hidden_dim, 1))
self.bz = np.zeros((hidden_dim, 1))
self.bh = np.zeros((hidden_dim, 1))
def _sigmoid(self, x):
return 1 / (1 + np.exp(-np.clip(x, -500, 500)))
def forward(self, inputs: List[np.ndarray]) -> List[np.ndarray]:
"""GRU前向传播"""
h_prev = np.zeros((self.hidden_dim, 1))
self.states = {'hs': {-1: h_prev}, 'gates': {}}
outputs = []
for t, x in enumerate(inputs):
concat = np.vstack((h_prev, x.reshape(-1, 1)))
# 重置门：控制历史信息的参与度
r_t = self._sigmoid(self.Wr @ concat + self.br)
# 更新门：控制历史状态与候选状态的混合
z_t = self._sigmoid(self.Wz @ concat + self.bz)
# 候选隐藏状态（受重置门调制）
concat_reset = np.vstack((r_t * h_prev, x.reshape(-1, 1)))
h_tilde = np.tanh(self.Wh @ concat_reset + self.bh)
# 隐藏状态更新（线性插值）
h_t = (1 - z_t) * h_prev + z_t * h_tilde
self.states['gates'][t] = {'r': r_t, 'z': z_t, 'h_tilde': h_tilde}
self.states['hs'][t] = h_t
outputs.append(h_t)
h_prev = h_t
return outputs
def demonstrate_gradient_dynamics():
"""展示不同架构的梯度动态特性"""
# 参数设置
input_dim = 10
hidden_dim = 64
seq_lengths = [10, 20, 30, 50, 100]
rnn_grad_norms = []
lstm_grad_norms = []
np.random.seed(42)
for seq_len in seq_lengths:
# 生成随机输入序列
inputs = [np.random.randn(input_dim) for _ in range(seq_len)]
targets = [np.random.randint(0, 10) for _ in range(seq_len)]
# Vanilla RNN
rnn = ManualRNN(input_dim, hidden_dim, 10)
rnn.forward(inputs)
grads = rnn.backward(targets, learning_rate=0.0) # 仅计算梯度，不更新
# 计算初始时间步的梯度范数（反映长程依赖能力）
rnn_grad_norms.append(np.mean(grads['grad_norms'][:5]))
# LSTM
lstm = ManualLSTM(input_dim, hidden_dim)
lstm.forward(inputs)
lstm_grads = lstm.analyze_gradient_flow(seq_len)
lstm_grad_norms.append(np.mean(lstm_grads[:5]))
# 可视化
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
# 图1：梯度范数随序列长度的衰减（谱分析视角）
ax1 = axes[0, 0]
ax1.semilogy(seq_lengths, rnn_grad_norms, 'o-', color='crimson',
linewidth=2, markersize=8, label='Vanilla RNN')
ax1.semilogy(seq_lengths, lstm_grad_norms, 's-', color='forestgreen',
linewidth=2, markersize=8, label='LSTM')
ax1.set_xlabel('Sequence Length', fontsize=11)
ax1.set_ylabel('Gradient Norm (log scale)', fontsize=11)
ax1.set_title('Spectral Analysis: Gradient Flow Stability', fontsize=12, fontweight='bold')
ax1.legend()
ax1.grid(True, alpha=0.3)
# 图2：LSTM门控激活可视化
ax2 = axes[0, 1]
lstm_demo = ManualLSTM(10, 32)
test_seq = [np.random.randn(10) for _ in range(20)]
lstm_demo.forward(test_seq)
gate_values = {'forget': [], 'input': [], 'output': []}
for t in range(20):
gates = lstm_demo.states['gates'][t]
gate_values['forget'].append(np.mean(gates['f']))
gate_values['input'].append(np.mean(gates['i']))
gate_values['output'].append(np.mean(gates['o']))
time_steps = range(20)
ax2.plot(time_steps, gate_values['forget'], label='Forget Gate',
color='steelblue', linewidth=2)
ax2.plot(time_steps, gate_values['input'], label='Input Gate',
color='darkorange', linewidth=2)
ax2.plot(time_steps, gate_values['output'], label='Output Gate',
color='forestgreen', linewidth=2)
ax2.set_xlabel('Time Step', fontsize=11)
ax2.set_ylabel('Gate Activation (Mean)', fontsize=11)
ax2.set_title('LSTM Gating Dynamics', fontsize=12, fontweight='bold')
ax2.legend()
ax2.grid(True, alpha=0.3)
# 图3：GRU门控对比（重置门 vs 更新门）
ax3 = axes[1, 0]
gru_demo = ManualGRU(10, 32)
gru_demo.forward(test_seq)
reset_vals = [np.mean(gru_demo.states['gates'][t]['r']) for t in range(20)]
update_vals = [np.mean(gru_demo.states['gates'][t]['z']) for t in range(20)]
ax3.plot(time_steps, reset_vals, label='Reset Gate',
color='purple', linewidth=2, linestyle='--')
ax3.plot(time_steps, update_vals, label='Update Gate',
color='teal', linewidth=2)
ax3.set_xlabel('Time Step', fontsize=11)
ax3.set_ylabel('Gate Activation (Mean)', fontsize=11)
ax3.set_title('GRU Simplified Gating Structure', fontsize=12, fontweight='bold')
ax3.legend()
ax3.grid(True, alpha=0.3)
# 图4：长期依赖任务中的梯度传播对比
ax4 = axes[1, 1]
long_seq_len = 100
long_inputs = [np.random.randn(10) for _ in range(long_seq_len)]
# 计算每个时间步的梯度范数
rnn_long = ManualRNN(10, 32, 10)
rnn_long.forward(long_inputs)
rnn_targets = [0] * long_seq_len
long_grads = rnn_long.backward(rnn_targets, learning_rate=0.0)
lstm_long = ManualLSTM(10, 32)
lstm_long.forward(long_inputs)
lstm_long_grads = lstm_long.analyze_gradient_flow(long_seq_len)
steps = range(long_seq_len)
ax4.semilogy(steps, long_grads['grad_norms'], alpha=0.6, color='crimson',
label='RNN Gradient Norm')
ax4.semilogy(steps, lstm_long_grads, alpha=0.6, color='forestgreen',
label='LSTM Cell State Gradient')
ax4.axhline(y=1e-7, color='red', linestyle=':', alpha=0.5, label='Vanishing Threshold')
ax4.set_xlabel('Backpropagation Time Step', fontsize=11)
ax4.set_ylabel('Gradient Magnitude', fontsize=11)
ax4.set_title('Long-Range Dependency: Gradient Propagation', fontsize=12, fontweight='bold')
ax4.legend()
ax4.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('rnn_gating_analysis.png', dpi=150, bbox_inches='tight')
print("Visualization saved to rnn_gating_analysis.png")
plt.show()
if __name__ == "__main__":
demonstrate_gradient_dynamics()10.1.2 Seq2Seq架构与教师强制的暴露偏差问题，计划采样的退火策略
序列到序列学习框架通过编码器将输入序列压缩为固定长度的上下文向量，解码器基于该向量自回归地生成输出序列。训练阶段采用教师强制策略，即解码器接收真实的前序标记作为下一步输入，而非自身生成的标记。这种训练与推理阶段输入分布的不一致性导致暴露偏差：模型在训练中从未接触自身生成的错误标记，推理时面对累积误差缺乏恢复能力，性能随序列长度增加而急剧退化。
计划采样作为一种课程学习策略，通过在训练过程中逐渐将解码器输入从真实标记过渡到模型生成标记来缓解暴露偏差。初始阶段以较大概率使用真实标记确保稳定学习，随着训练进行线性或指数退火降低真实标记使用概率，使模型逐步适应自身生成的分布。这种退火机制模拟了课程学习的渐进难度提升，在保持训练稳定性的同时增强模型的自洽生成能力。
脚本：seq2seq_exposure_bias.py
Python
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脚本内容：Seq2Seq架构、Teacher Forcing与Scheduled Sampling实现
使用方式：直接运行 python seq2seq_exposure_bias.py
依赖包：torch, numpy, matplotlib
功能说明：
1. 实现完整的Seq2Seq架构（Encoder-Decoder）
2. 对比Teacher Forcing与Scheduled Sampling的训练动态
3. 可视化暴露偏差（Exposure Bias）现象
4. 展示不同退火策略（线性/指数/反Sigmoid）的效果
"""
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple, Optional
import random
class Encoder(nn.Module):
"""标准GRU编码器，支持多层堆叠"""
def __init__(self, vocab_size: int, embed_dim: int, hidden_dim: int,
num_layers: int = 2, dropout: float = 0.1):
super().__init__()
self.hidden_dim = hidden_dim
self.num_layers = num_layers
self.embedding = nn.Embedding(vocab_size, embed_dim)
self.gru = nn.GRU(embed_dim, hidden_dim, num_layers,
batch_first=True, dropout=dropout if num_layers > 1 else 0)
def forward(self, src: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Args:
src: [batch_size, src_len]
Returns:
outputs: [batch_size, src_len, hidden_dim]
hidden: [num_layers, batch_size, hidden_dim]
"""
embedded = self.embedding(src) # [batch, src_len, embed_dim]
outputs, hidden = self.gru(embedded)
return outputs, hidden
class AttentionDecoder(nn.Module):
"""带Bahdanau注意力的解码器，支持Scheduled Sampling"""
def __init__(self, vocab_size: int, embed_dim: int, hidden_dim: int,
num_layers: int = 2, dropout: float = 0.1):
super().__init__()
self.vocab_size = vocab_size
self.hidden_dim = hidden_dim
self.num_layers = num_layers
self.embedding = nn.Embedding(vocab_size, embed_dim)
# 注意力机制参数
self.attn_W = nn.Linear(hidden_dim, hidden_dim) # 编码器输出变换
self.attn_U = nn.Linear(hidden_dim, hidden_dim) # 解码器隐藏状态变换
self.attn_v = nn.Linear(hidden_dim, 1) # 注意力评分
# GRU输入维度为embedding_dim + hidden_dim（上下文向量）
self.gru = nn.GRU(embed_dim + hidden_dim, hidden_dim, num_layers,
batch_first=True, dropout=dropout if num_layers > 1 else 0)
# 输出投影
self.fc_out = nn.Linear(hidden_dim, vocab_size)
self.dropout = nn.Dropout(dropout)
def compute_attention(self, hidden: torch.Tensor,
encoder_outputs: torch.Tensor) -> torch.Tensor:
"""
Bahdanau注意力计算
hidden: [num_layers, batch, hidden_dim] -> 取最后一层
encoder_outputs: [batch, src_len, hidden_dim]
"""
# 取最后一层隐藏状态 [batch, hidden_dim]
hidden_last = hidden[-1].unsqueeze(1) # [batch, 1, hidden_dim]
# 计算注意力分数
# [batch, src_len, hidden_dim]
attn_scores = self.attn_v(torch.tanh(self.attn_W(encoder_outputs) +
self.attn_U(hidden_last)))
# [batch, src_len, 1] -> [batch, src_len]
attn_scores = attn_scores.squeeze(-1)
attn_weights = torch.softmax(attn_scores, dim=1)
# 计算上下文向量 [batch, 1, hidden_dim]
context = torch.bmm(attn_weights.unsqueeze(1), encoder_outputs)
return context.squeeze(1), attn_weights
def forward_step(self, input_token: torch.Tensor, hidden: torch.Tensor,
encoder_outputs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""单步解码"""
embedded = self.dropout(self.embedding(input_token)) # [batch, embed_dim]
# 计算注意力上下文
context, attn_weights = self.compute_attention(hidden, encoder_outputs)
# 拼接embedding和上下文
gru_input = torch.cat([embedded, context], dim=1).unsqueeze(1)
# GRU前向
output, hidden = self.gru(gru_input, hidden)
output = output.squeeze(1) # [batch, hidden_dim]
# 预测下一个token
prediction = self.fc_out(output) # [batch, vocab_size]
return prediction, hidden, attn_weights
def forward(self, trg: torch.Tensor, hidden: torch.Tensor,
encoder_outputs: torch.Tensor,
teacher_forcing_ratio: float = 0.5) -> Tuple[torch.Tensor, torch.Tensor]:
"""
前向传播，支持Scheduled Sampling
Args:
trg: [batch_size, trg_len] 目标序列
hidden: [num_layers, batch_size, hidden_dim] 编码器最终隐藏状态
encoder_outputs: [batch_size, src_len, hidden_dim]
teacher_forcing_ratio: 使用真实标签的概率（退火策略控制）
"""
batch_size = trg.size(0)
trg_len = trg.size(1)
outputs = torch.zeros(batch_size, trg_len, self.vocab_size).to(trg.device)
attention_weights = []
# 初始输入为<BOS>标记（假设为0）
input_token = torch.zeros(batch_size, dtype=torch.long).to(trg.device)
for t in range(trg_len):
prediction, hidden, attn = self.forward_step(input_token, hidden, encoder_outputs)
outputs[:, t, :] = prediction
attention_weights.append(attn)
# Scheduled Sampling决策
teacher_force = random.random() < teacher_forcing_ratio
if teacher_force:
input_token = trg[:, t] # 使用真实标记
else:
input_token = prediction.argmax(1) # 使用模型预测
return outputs, torch.stack(attention_weights, dim=1)
class Seq2Seq(nn.Module):
"""完整的Seq2Seq模型"""
def __init__(self, encoder: Encoder, decoder: AttentionDecoder, device: torch.device):
super().__init__()
self.encoder = encoder
self.decoder = decoder
self.device = device
def forward(self, src: torch.Tensor, trg: torch.Tensor,
teacher_forcing_ratio: float = 0.5):
encoder_outputs, hidden = self.encoder(src)
outputs, attentions = self.decoder(trg, hidden, encoder_outputs,
teacher_forcing_ratio)
return outputs, attentions
class ScheduledSamplingTrainer:
"""实现不同退火策略的训练器"""
def __init__(self, model: Seq2Seq, strategy: str = 'linear',
start_ratio: float = 1.0, end_ratio: float = 0.0):
self.model = model
self.strategy = strategy
self.start_ratio = start_ratio
self.end_ratio = end_ratio
def get_teacher_forcing_ratio(self, epoch: int, total_epochs: int) -> float:
"""计算当前epoch的教师强制比例"""
progress = epoch / total_epochs
if self.strategy == 'linear':
# 线性退火
return self.start_ratio - (self.start_ratio - self.end_ratio) * progress
elif self.strategy == 'exponential':
# 指数退火
decay_rate = 5.0 # 控制衰减速度
return self.end_ratio + (self.start_ratio - self.end_ratio) * \
np.exp(-decay_rate * progress)
elif self.strategy == 'inverse_sigmoid':
# 反Sigmoid退火（先慢后快）
k = 10.0
return self.end_ratio + (self.start_ratio - self.end_ratio) * \
(k / (k + np.exp(k * progress)))
else:
return self.start_ratio
def train_epoch(self, dataloader, optimizer, criterion, epoch: int,
total_epochs: int):
"""训练单个epoch"""
self.model.train()
total_loss = 0
teacher_ratio = self.get_teacher_forcing_ratio(epoch, total_epochs)
for batch_idx, (src, trg) in enumerate(dataloader):
src, trg = src.to(self.model.device), trg.to(self.model.device)
optimizer.zero_grad()
output, _ = self.model(src, trg, teacher_forcing_ratio=teacher_ratio)
# 计算损失（忽略padding）
output_dim = output.shape[-1]
output = output[:, 1:].reshape(-1, output_dim) # 跳过<BOS>
trg = trg[:, 1:].reshape(-1)
loss = criterion(output, trg)
loss.backward()
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
optimizer.step()
total_loss += loss.item()
return total_loss / len(dataloader), teacher_ratio
def simulate_exposure_bias_experiment():
"""
模拟暴露偏差实验：
对比纯Teacher Forcing、Scheduled Sampling与Oracle（使用真实前序）的性能差距
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 模拟数据参数
vocab_size = 1000
batch_size = 32
src_len = 20
trg_len = 20
hidden_dim = 256
epochs = 50
# 生成模拟数据（随机序列复制任务）
def generate_dummy_data(num_samples=1000):
src = torch.randint(1, vocab_size, (num_samples, src_len))
# 目标序列是源序列的某种变换（这里简化为复制）
trg = torch.cat([torch.zeros(num_samples, 1, dtype=torch.long),
src[:, :trg_len-1]], dim=1)
return torch.utils.data.DataLoader(
list(zip(src, trg)), batch_size=batch_size, shuffle=True
)
train_loader = generate_dummy_data()
# 三种训练策略对比
strategies = {
'Teacher Forcing (ratio=1.0)': lambda e: 1.0,
'Scheduled Sampling (linear)': lambda e: max(0.0, 1.0 - e/epochs),
'Free Running (ratio=0.0)': lambda e: 0.0
}
results = {name: {'losses': [], 'exposure_bias_metric': []}
for name in strategies.keys()}
for strategy_name, ratio_fn in strategies.items():
print(f"\nTraining with {strategy_name}...")
# 重新初始化模型
encoder = Encoder(vocab_size, 128, hidden_dim, num_layers=2)
decoder = AttentionDecoder(vocab_size, 128, hidden_dim, num_layers=2)
model = Seq2Seq(encoder, decoder, device).to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss(ignore_index=0)
epoch_losses = []
exposure_scores = [] # 衡量暴露偏差：测试时使用自生成输入的性能下降
for epoch in range(epochs):
model.train()
epoch_loss = 0
teacher_ratio = ratio_fn(epoch)
for src, trg in train_loader:
src, trg = src.to(device), trg.to(device)
optimizer.zero_grad()
output, _ = model(src, trg, teacher_forcing_ratio=teacher_ratio)
output_dim = output.shape[-1]
loss = criterion(output[:, 1:].reshape(-1, output_dim),
trg[:, 1:].reshape(-1))
loss.backward()
optimizer.step()
epoch_loss += loss.item()
avg_loss = epoch_loss / len(train_loader)
epoch_losses.append(avg_loss)
# 计算暴露偏差：训练损失 vs 推理损失（无teacher forcing）
model.eval()
with torch.no_grad():
val_src, val_trg = next(iter(train_loader))
val_src, val_trg = val_src.to(device), val_trg.to(device)
# 教师强制下的损失（上界）
tf_output, _ = model(val_src, val_trg, teacher_forcing_ratio=1.0)
tf_loss = criterion(tf_output[:, 1:].reshape(-1, output_dim),
val_trg[:, 1:].reshape(-1)).item()
# 自由运行下的损失（真实推理场景）
fr_output, _ = model(val_src, val_trg, teacher_forcing_ratio=0.0)
fr_loss = criterion(fr_output[:, 1:].reshape(-1, output_dim),
val_trg[:, 1:].reshape(-1)).item()
# 暴露偏差 = 自由运行损失 - 教师强制损失
exposure_scores.append(fr_loss - tf_loss)
results[strategy_name]['losses'] = epoch_losses
results[strategy_name]['exposure_bias_metric'] = exposure_scores
# 可视化结果
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
# 图1：训练损失曲线
ax1 = axes[0, 0]
colors = ['crimson', 'forestgreen', 'steelblue']
for (name, data), color in zip(results.items(), colors):
ax1.plot(data['losses'], label=name, color=color, linewidth=2)
ax1.set_xlabel('Epoch', fontsize=11)
ax1.set_ylabel('Training Loss', fontsize=11)
ax1.set_title('Training Dynamics: Loss Convergence', fontsize=12, fontweight='bold')
ax1.legend()
ax1.grid(True, alpha=0.3)
# 图2：暴露偏差随训练变化
ax2 = axes[0, 1]
for (name, data), color in zip(results.items(), colors):
ax2.plot(data['exposure_bias_metric'], label=name, color=color, linewidth=2)
ax2.axhline(y=0, color='black', linestyle='--', alpha=0.3)
ax2.set_xlabel('Epoch', fontsize=11)
ax2.set_ylabel('Exposure Bias (Free-Run Loss - TF Loss)', fontsize=11)
ax2.set_title('Exposure Bias Accumulation', fontsize=12, fontweight='bold')
ax2.legend()
ax2.grid(True, alpha=0.3)
# 图3：退火策略曲线对比
ax3 = axes[1, 0]
epochs_range = range(epochs)
linear = [max(0.0, 1.0 - e/epochs) for e in epochs_range]
exponential = [np.exp(-5.0 * e/epochs) for e in epochs_range]
inverse_sigmoid = [10.0 / (10.0 + np.exp(10.0 * e/epochs)) for e in epochs_range]
ax3.plot(epochs_range, linear, label='Linear', linewidth=2, color='darkorange')
ax3.plot(epochs_range, exponential, label='Exponential', linewidth=2, color='purple')
ax3.plot(epochs_range, inverse_sigmoid, label='Inverse Sigmoid', linewidth=2, color='teal')
ax3.set_xlabel('Epoch', fontsize=11)
ax3.set_ylabel('Teacher Forcing Ratio', fontsize=11)
ax3.set_title('Scheduled Sampling: Annealing Strategies', fontsize=12, fontweight='bold')
ax3.legend()
ax3.grid(True, alpha=0.3)
# 图4：注意力对齐可视化（Scheduled Sampling训练后的模型）
ax4 = axes[1, 1]
# 使用最后一个训练好的模型（Scheduled Sampling）
model.eval()
with torch.no_grad():
test_src = torch.randint(1, vocab_size, (1, src_len)).to(device)
test_trg = torch.zeros(1, trg_len, dtype=torch.long).to(device)
encoder_outs, hidden = model.encoder(test_src)
_, attentions = model.decoder(test_trg, hidden, encoder_outs, 0.0)
# 转换为numpy热力图
attn_map = attentions[0].cpu().numpy() # [trg_len, src_len]
im = ax4.imshow(attn_map, cmap='viridis', aspect='auto')
ax4.set_xlabel('Source Sequence Position', fontsize=11)
ax4.set_ylabel('Target Generation Step', fontsize=11)
ax4.set_title('Attention Alignment (Post-Scheduled Sampling)',
fontsize=12, fontweight='bold')
plt.colorbar(im, ax=ax4, label='Attention Weight')
plt.tight_layout()
plt.savefig('seq2seq_exposure_bias.png', dpi=150, bbox_inches='tight')
print("\nVisualization saved to seq2seq_exposure_bias.png")
plt.show()
if __name__ == "__main__":
simulate_exposure_bias_experiment()10.1.3 双向RNN与多层堆叠的并行化实现，CuDNN内核优化原理
双向循环神经网络通过同时维护前向和后向两个隐藏状态序列，捕获输入序列的双向上下文依赖。前向层按时间正序处理序列，后向层按时间逆序处理，二者的隐藏状态在输出层拼接或累加，使每个时间步的表示融合了过去与未来的信息。这种架构显著提升了序列标注和句子编码任务的性能，但将计算依赖翻倍，对并行化提出了更高要求。
多层堆叠通过垂直方向叠加循环层构建深层时序模型，下层隐藏状态作为上层输入，逐级提取抽象特征。层间可以添加残差连接或密集连接以促进梯度流动。深层双向架构的计算图呈现二维依赖结构，既包含水平方向的时间依赖，又包含垂直方向的层间依赖。
CuDNN库通过内核融合与并行化策略优化循环层计算。首先，将矩阵乘法、偏置加法和激活函数融合为单一内核，减少显存带宽消耗；其次，采用层间融合技术，将多层双向RNN的单元计算流水线化，隐藏层间同步开销；最后，针对长序列采用时间分块策略，在保持循环依赖的同时利用GPU的流式多处理器并行处理不同时间块。这些优化使得深度双向LSTM在GPU上的吞吐量相比朴素实现提升十倍以上。
脚本：birnn_parallelization.py
Python
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脚本内容：双向RNN、多层堆叠与并行化优化实现
使用方式：直接运行 python birnn_parallelization.py
依赖包：torch, numpy, matplotlib, time
功能说明：
1. 实现手动双向RNN与PyTorch CuDNN优化版本的对比
2. 展示层间融合与时间分块的优化效果
3. 可视化多层双向RNN的特征提取层次
4. 分析不同并行策略的计算效率
"""
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple, Optional
import time
class ManualBiRNN:
"""手动实现的双向RNN，用于教学目的"""
def __init__(self, input_dim: int, hidden_dim: int, num_layers: int = 1):
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.num_layers = num_layers
# 为每层创建前向和后向参数
self.forward_params = []
self.backward_params = []
for layer in range(num_layers):
layer_input_dim = input_dim if layer == 0 else hidden_dim * 2
# 前向参数
fwd = {
'W_ih': torch.randn(hidden_dim, layer_input_dim) * 0.01,
'W_hh': torch.eye(hidden_dim) * 0.9, # 正交初始化促进梯度流
'b_ih': torch.zeros(hidden_dim),
'b_hh': torch.zeros(hidden_dim)
}
# 后向参数
bwd = {
'W_ih': torch.randn(hidden_dim, layer_input_dim) * 0.01,
'W_hh': torch.eye(hidden_dim) * 0.9,
'b_ih': torch.zeros(hidden_dim),
'b_hh': torch.zeros(hidden_dim)
}
self.forward_params.append(fwd)
self.backward_params.append(bwd)
def _lstm_cell(self, x: torch.Tensor, h: torch.Tensor, c: torch.Tensor,
params: dict) -> Tuple[torch.Tensor, torch.Tensor]:
"""单层LSTM单元计算"""
# 合并输入和隐藏状态
combined = torch.cat([x, h], dim=-1)
# 门控计算
gates = torch.mm(combined, torch.cat([params['W_ih'], params['W_hh']], dim=1).t())
gates = gates + params['b_ih'] + params['b_hh']
i, f, g, o = gates.chunk(4, 1)
i = torch.sigmoid(i)
f = torch.sigmoid(f)
g = torch.tanh(g)
o = torch.sigmoid(o)
c_new = f * c + i * g
h_new = o * torch.tanh(c_new)
return h_new, c_new
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
手动前向传播，展示双向和多层机制
x: [batch, seq_len, input_dim]
"""
batch_size, seq_len, _ = x.size()
device = x.device
current_input = x
for layer in range(self.num_layers):
# 前向传播
h_fwd = torch.zeros(batch_size, self.hidden_dim).to(device)
c_fwd = torch.zeros(batch_size, self.hidden_dim).to(device)
fwd_outputs = []
for t in range(seq_len):
h_fwd, c_fwd = self._lstm_cell(
current_input[:, t, :], h_fwd, c_fwd,
self.forward_params[layer]
)
fwd_outputs.append(h_fwd)
fwd_outputs = torch.stack(fwd_outputs, dim=1) # [batch, seq, hidden]
# 后向传播
h_bwd = torch.zeros(batch_size, self.hidden_dim).to(device)
c_bwd = torch.zeros(batch_size, self.hidden_dim).to(device)
bwd_outputs = []
for t in reversed(range(seq_len)):
h_bwd, c_bwd = self._lstm_cell(
current_input[:, t, :], h_bwd, c_bwd,
self.backward_params[layer]
)
bwd_outputs.insert(0, h_bwd)
bwd_outputs = torch.stack(bwd_outputs, dim=1)
# 拼接前向和后向
layer_output = torch.cat([fwd_outputs, bwd_outputs], dim=-1)
current_input = layer_output # 作为下一层输入
return current_input
class OptimizedDeepBiRNN(nn.Module):
"""使用PyTorch CuDNN优化的高效深度双向RNN"""
def __init__(self, input_dim: int, hidden_dim: int, num_layers: int = 3,
dropout: float = 0.3, rnn_type: str = 'LSTM'):
super().__init__()
self.rnn_type = rnn_type
rnn_cls = nn.LSTM if rnn_type == 'LSTM' else nn.GRU
# CuDNN优化要求：batch_first=True, 参数连续存储
self.rnn = rnn_cls(
input_size=input_dim,
hidden_size=hidden_dim,
num_layers=num_layers,
bidirectional=True,
batch_first=True,
dropout=dropout if num_layers > 1 else 0,
)
# 层归一化（Pre-LN风格）稳定深层训练
self.layer_norms = nn.ModuleList([
nn.LayerNorm(hidden_dim * 2) for _ in range(num_layers)
])
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
返回输出和层间特征用于可视化
"""
# CuDNN内核自动处理层间融合和时间并行
outputs, (hidden, cell) = self.rnn(x)
# 分层提取特征用于分析
layer_features = []
current = x
# 手动逐层前向以获取中间表示（分析用，非生产环境使用）
if self.rnn_type == 'LSTM':
# 使用权重展开分析
weight_ih_l0 = self.rnn.weight_ih_l0
layer_features.append(outputs)
return outputs, hidden
class ParallelizationAnalyzer:
"""分析不同并行化策略的效率"""
def __init__(self):
self.results = {}
def benchmark_implementation(self, model, input_tensor, num_runs=100):
"""基准测试模型推理速度"""
device = next(model.parameters()).device if hasattr(model, 'parameters') else torch.device('cpu')
input_tensor = input_tensor.to(device)
# 预热
for _ in range(10):
_ = model(input_tensor)
if torch.cuda.is_available():
torch.cuda.synchronize()
start = time.time()
for _ in range(num_runs):
output = model(input_tensor)
if torch.cuda.is_available():
torch.cuda.synchronize()
end = time.time()
avg_time = (end - start) / num_runs * 1000 # ms
return avg_time
def compare_strategies(self, input_dims: List[int], seq_lengths: List[int],
hidden_dim: int = 256, num_layers: int = 3):
"""对比不同配置下的性能"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
results = {
'manual': [],
'cudnn_optimized': [],
'config': []
}
for input_dim in input_dims:
for seq_len in seq_lengths:
batch_size = 32
# 创建输入
x = torch.randn(batch_size, seq_len, input_dim)
# 手动实现（仅小规模测试）
if seq_len <= 50 and input_dim <= 64:
manual_model = ManualBiRNN(input_dim, hidden_dim//2, num_layers)
manual_time = self.benchmark_implementation(manual_model, x, num_runs=10)
results['manual'].append(manual_time)
else:
results['manual'].append(None)
# CuDNN优化版本
cudnn_model = OptimizedDeepBiRNN(input_dim, hidden_dim//2, num_layers).to(device)
cudnn_time = self.benchmark_implementation(cudnn_model, x, num_runs=100)
results['cudnn_optimized'].append(cudnn_time)
results['config'].append(f"dim={input_dim}, len={seq_len}")
return results
def visualize_layer_activations(self, model: OptimizedDeepBiRNN,
sample_input: torch.Tensor):
"""可视化多层双向RNN的层间特征演化"""
model.eval()
with torch.no_grad():
outputs, hidden = model(sample_input)
# 分析不同层的梯度流和特征分布
batch_size, seq_len, feat_dim = outputs.size()
# 分离不同双向方向的特征
fwd_feat = outputs[:, :, :feat_dim//2]
bwd_feat = outputs[:, :, feat_dim//2:]
return fwd_feat.numpy(), bwd_feat.numpy()
def create_visualization():
"""创建完整的可视化分析"""
# 设置参数
input_dim = 64
hidden_dim = 128
num_layers = 3
seq_len = 100
batch_size = 16
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 生成模拟序列数据（正弦波叠加噪声，模拟时序信号）
t = np.linspace(0, 4*np.pi, seq_len)
signal = np.sin(t) + 0.5 * np.sin(3*t) + 0.3 * np.random.randn(seq_len)
x = torch.FloatTensor(signal).view(1, seq_len, 1).repeat(batch_size, 1, input_dim)
fig, axes = plt.subplots(2, 2, figsize=(15, 11))
# 图1：手动 vs CuDNN 性能对比
ax1 = axes[0, 0]
analyzer = ParallelizationAnalyzer()
seq_lengths = [10, 20, 50, 100, 200]
manual_times = []
cudnn_times = []
for sl in seq_lengths:
x_test = torch.randn(32, sl, input_dim)
# 手动版本（小规模）
if sl <= 50:
manual_model = ManualBiRNN(input_dim, hidden_dim//2, 2)
t_manual = analyzer.benchmark_implementation(manual_model, x_test, 10)
manual_times.append(t_manual)
else:
manual_times.append(np.nan)
# CuDNN版本
cudnn_model = OptimizedDeepBiRNN(input_dim, hidden_dim//2, 2).to(device)
t_cudnn = analyzer.benchmark_implementation(cudnn_model, x_test, 100)
cudnn_times.append(t_cudnn)
ax1.plot(seq_lengths, manual_times, 'o-', label='Manual Implementation',
color='crimson', linewidth=2, markersize=8)
ax1.plot(seq_lengths, cudnn_times, 's-', label='CuDNN Optimized',
color='forestgreen', linewidth=2, markersize=8)
ax1.set_xlabel('Sequence Length', fontsize=11)
ax1.set_ylabel('Inference Time (ms)', fontsize=11)
ax1.set_title('Parallelization Efficiency: Manual vs CuDNN',
fontsize=12, fontweight='bold')
ax1.legend()
ax1.grid(True, alpha=0.3)
ax1.set_yscale('log')
# 图2：双向特征可视化（热力图）
ax2 = axes[0, 1]
model = OptimizedDeepBiRNN(input_dim, hidden_dim//2, num_layers).to(device)
model.eval()
with torch.no_grad():
x_device = x.to(device)
outputs, _ = model(x_device)
# 取第一个样本，展示前向和后向特征的差异
sample_output = outputs[0].cpu().numpy() # [seq_len, hidden*2]
# 计算前向和后向特征的差异度
fwd = sample_output[:, :hidden_dim//2]
bwd = sample_output[:, hidden_dim//2:]
im = ax2.imshow(sample_output.T, aspect='auto', cmap='RdBu_r',
vmin=-2, vmax=2, interpolation='bilinear')
ax2.set_xlabel('Time Step', fontsize=11)
ax2.set_ylabel('Hidden Units (Fwd | Bwd)', fontsize=11)
ax2.set_title('Bidirectional Feature Maps (Layer 3)',
fontsize=12, fontweight='bold')
ax2.axhline(y=hidden_dim//2, color='black', linestyle='--', linewidth=2)
ax2.text(seq_len*0.02, hidden_dim//4, 'Forward', fontsize=10,
color='white', fontweight='bold')
ax2.text(seq_len*0.02, hidden_dim*0.75, 'Backward', fontsize=10,
color='white', fontweight='bold')
plt.colorbar(im, ax=ax2, label='Activation')
# 图3：多层特征的时间演化（不同层的时间常数）
ax3 = axes[1, 0]
# 分析不同层的"记忆长度"（自相关分析）
layer_memories = []
for layer_idx in [0, 1, 2]:
# 模拟不同层的特征（实际应从模型中提取）
# 高层应具有更长的时间常数
time_const = 5 * (layer_idx + 1)
memory_curve = np.exp(-np.arange(50) / time_const)
layer_memories.append(memory_curve)
for idx, curve in enumerate(layer_memories):
ax3.plot(curve, label=f'Layer {idx+1} (τ={5*(idx+1)})', linewidth=2)
ax3.set_xlabel('Time Lag', fontsize=11)
ax3.set_ylabel('Auto-correlation', fontsize=11)
ax3.set_title('Hierarchical Temporal Receptive Fields',
fontsize=12, fontweight='bold')
ax3.legend()
ax3.grid(True, alpha=0.3)
# 图4：层间融合与分块策略的吞吐量分析
ax4 = axes[1, 1]
# 模拟不同优化策略的吞吐量
batch_sizes = [1, 8, 16, 32, 64, 128]
# 朴素实现：受限于内存带宽和顺序计算
naive_throughput = [10 * bs for bs in batch_sizes]
naive_throughput = [min(v, 500) for v in naive_throughput] # 瓶颈
# 层融合：减少kernel launch开销
fused_throughput = [15 * bs for bs in batch_sizes]
fused_throughput = [min(v, 800) for v in fused_throughput]
# 时间分块+并行：利用GPU多流处理器
tiled_throughput = [25 * bs for bs in batch_sizes]
tiled_throughput = [min(v, 1200 + 5*bs) for bs, v in zip(batch_sizes, tiled_throughput)]
ax4.plot(batch_sizes, naive_throughput, 'o--', label='Naive Implementation',
color='crimson', linewidth=2)
ax4.plot(batch_sizes, fused_throughput, 's-', label='Layer Fusion (CuDNN)',
color='darkorange', linewidth=2)
ax4.plot(batch_sizes, tiled_throughput, '^-', label='Time Tiling + Parallel',
color='forestgreen', linewidth=2)
ax4.set_xlabel('Batch Size', fontsize=11)
ax4.set_ylabel('Throughput (samples/sec)', fontsize=11)
ax4.set_title('GPU Kernel Optimization Strategies',
fontsize=12, fontweight='bold')
ax4.legend()
ax4.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('birnn_parallelization.png', dpi=150, bbox_inches='tight')
print("Visualization saved to birnn_parallelization.png")
plt.show()
if __name__ == "__main__":
create_visualization()10.2 注意力机制与Transformer架构
10.2.1 Bahdanau注意力与Luong注意力的分数函数对比，键-查询-值矩阵乘法视角
注意力机制通过动态权重分配机制，使模型能够在生成每个输出元素时聚焦于输入序列的相关部分。Bahdanau注意力采用加法评分函数，将解码器当前隐藏状态与编码器各时间步输出通过独立线性变换投影至共同空间后相加，经非线性激活映射为标量分数。这种机制允许查询与键在不同表示空间中交互，增强了模型的表达能力，但增加了计算复杂度。
Luong注意力提出乘法评分函数，直接使用解码器隐藏状态与编码器输出的点积或经参数矩阵变换后的点积计算相似度。点积机制假设查询与键处于兼容的表示空间，通过矩阵乘法直接衡量向量夹角，计算效率显著高于加法机制。泛化点积引入可学习的权重矩阵调节注意力分布的几何特性。
键查询值矩阵乘法视角将注意力重新诠释为寻址操作。键向量构建可搜索的内容地址，查询向量作为搜索请求，通过相似度计算定位相关记忆，值向量提供实际存储的内容信息。这种抽象使得注意力机制脱离序列建模的特定语境，成为通用的关联记忆检索原语，为后续的自注意力架构奠定了概念基础。
脚本：attention_mechanisms_comparison.py
Python
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脚本内容：Bahdanau与Luong注意力机制对比实现
使用方式：直接运行 python attention_mechanisms_comparison.py
依赖包：torch, numpy, matplotlib
功能说明：
1. 手动实现Bahdanau（Additive）和Luong（Dot/General/Concat）注意力
2. 展示键-查询-值（K-Q-V）矩阵乘法的寻址视角
3. 可视化不同评分函数的对齐矩阵（Alignment Matrix）
4. 对比计算效率与表达能力
"""
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Optional
class BahdanauAttention(nn.Module):
"""Bahdanau（Additive）注意力实现"""
def __init__(self, hidden_dim: int, encoder_dim: int):
super().__init__()
self.hidden_dim = hidden_dim
self.encoder_dim = encoder_dim
# 独立线性变换：将解码器状态与编码器输出映射到相同空间
self.W_query = nn.Linear(hidden_dim, hidden_dim, bias=False)
self.W_key = nn.Linear(encoder_dim, hidden_dim, bias=False)
self.v = nn.Linear(hidden_dim, 1, bias=False) # 评分向量
self.tanh = nn.Tanh()
def forward(self, query: torch.Tensor, keys: torch.Tensor,
values: torch.Tensor, mask: Optional[torch.Tensor] = None):
"""
Args:
query: [batch, hidden_dim] - 解码器当前状态
keys: [batch, src_len, encoder_dim] - 编码器所有输出（作为地址）
values: [batch, src_len, value_dim] - 编码器所有输出（作为内容）
mask: [batch, src_len] - 掩码
Returns:
context: [batch, value_dim] - 加权上下文向量
weights: [batch, src_len] - 注意力分布
scores: [batch, src_len] - 原始分数（用于可视化）
"""
batch_size, src_len, _ = keys.size()
# 变换查询和键到共同空间（Additive机制核心）
# query: [batch, hidden] -> [batch, 1, hidden] -> [batch, src_len, hidden]
query_transformed = self.W_query(query).unsqueeze(1).expand(-1, src_len, -1)
# keys: [batch, src_len, encoder_dim] -> [batch, src_len, hidden]
keys_transformed = self.W_key(keys)
# 相加后非线性激活（允许不同空间的交互）
additive_scores = self.v(self.tanh(query_transformed + keys_transformed)).squeeze(-1)
# 应用掩码
if mask is not None:
additive_scores.masked_fill_(mask == 0, float('-inf'))
# 计算注意力权重
weights = torch.softmax(additive_scores, dim=-1) # [batch, src_len]
# 加权求和获取上下文（寻址操作）
context = torch.bmm(weights.unsqueeze(1), values).squeeze(1)
return context, weights, additive_scores
class LuongAttention(nn.Module):
"""Luong（Multiplicative）注意力实现，支持多种评分函数"""
def __init__(self, hidden_dim: int, encoder_dim: int, score_mode: str = 'general'):
super().__init__()
self.hidden_dim = hidden_dim
self.encoder_dim = encoder_dim
self.score_mode = score_mode
if score_mode == 'general':
# 可学习的参数矩阵调节注意力几何
self.W = nn.Linear(hidden_dim, encoder_dim, bias=False)
elif score_mode == 'concat':
# 拼接后变换（介于additive和dot之间）
self.W = nn.Linear(hidden_dim + encoder_dim, hidden_dim, bias=False)
self.v = nn.Linear(hidden_dim, 1, bias=False)
def forward(self, query: torch.Tensor, keys: torch.Tensor,
values: torch.Tensor, mask: Optional[torch.Tensor] = None):
"""
Args同BahdanauAttention
"""
batch_size, src_len, _ = keys.size()
if self.score_mode == 'dot':
# 假设查询与键空间兼容，直接点积
# query: [batch, hidden], keys: [batch, src_len, hidden]
# 要求 hidden_dim == encoder_dim
scores = torch.bmm(keys, query.unsqueeze(-1)).squeeze(-1)
elif self.score_mode == 'general':
# 通过参数矩阵调节后的点积
query_transformed = self.W(query) # [batch, encoder_dim]
scores = torch.bmm(keys, query_transformed.unsqueeze(-1)).squeeze(-1)
elif self.score_mode == 'concat':
# 扩展query与keys拼接
query_expanded = query.unsqueeze(1).expand(-1, src_len, -1)
concat = torch.cat([query_expanded, keys], dim=-1)
scores = self.v(torch.tanh(self.W(concat))).squeeze(-1)
# 掩码处理
if mask is not None:
scores.masked_fill_(mask == 0, float('-inf'))
weights = torch.softmax(scores, dim=-1)
context = torch.bmm(weights.unsqueeze(1), values).squeeze(1)
return context, weights, scores
class KeyQueryValuePerspective(nn.Module):
"""展示K-Q-V矩阵乘法视角的注意力实现"""
def __init__(self, dim: int, num_heads: int = 1):
super().__init__()
self.dim = dim
self.num_heads = num_heads
self.head_dim = dim // num_heads
# 将输入投影为K、Q、V（内容寻址机制）
self.W_k = nn.Linear(dim, dim, bias=False)
self.W_q = nn.Linear(dim, dim, bias=False)
self.W_v = nn.Linear(dim, dim, bias=False)
def forward(self, x: torch.Tensor):
"""
展示如何从输入提取K-Q-V并进行寻址
x: [batch, seq_len, dim]
"""
batch_size, seq_len, _ = x.size()
# 投影到K、Q、V空间（地址与内容的分离）
K = self.W_k(x) # 键：用于寻址的"地址标签"
Q = self.W_q(x) # 查询："检索请求"
V = self.W_v(x) # 值：实际存储的"内容"
# 计算注意力分数：查询与地址的匹配程度
scores = torch.bmm(Q, K.transpose(1, 2)) / np.sqrt(self.dim)
# 归一化得到寻址权重
attn_weights = torch.softmax(scores, dim=-1)
# 根据权重聚合内容
output = torch.bmm(attn_weights, V)
return output, attn_weights, (K, Q, V)
def visualize_attention_mechanisms():
"""可视化不同注意力机制的对比"""
# 创建模拟数据：机器翻译场景
batch_size = 4
src_len = 20
tgt_len = 15
hidden_dim = 256
encoder_dim = 512
# 模拟编码器输出（源语言）
encoder_outputs = torch.randn(batch_size, src_len, encoder_dim)
# 模拟解码器状态（目标语言生成中）
decoder_hidden = torch.randn(batch_size, hidden_dim)
# 模拟目标序列（用于展示对齐）
target_positions = torch.arange(tgt_len).float().unsqueeze(0).expand(batch_size, -1)
fig, axes = plt.subplots(2, 3, figsize=(16, 10))
# 1. Bahdanau注意力对齐矩阵
bahdanau = BahdanauAttention(hidden_dim, encoder_dim)
bahdanau.eval()
alignment_matrix = []
with torch.no_grad():
for t in range(tgt_len):
# 模拟解码器状态随时间变化（累加正弦位置编码）
query = decoder_hidden + 0.1 * torch.sin(target_positions[:, t:t+1] * 0.5).expand(-1, hidden_dim)
context, weights, scores = bahdanau(query, encoder_outputs, encoder_outputs)
alignment_matrix.append(weights[0].numpy()) # 取第一个样本
alignment_matrix = np.array(alignment_matrix)
ax1 = axes[0, 0]
im1 = ax1.imshow(alignment_matrix, cmap='viridis', aspect='auto', interpolation='nearest')
ax1.set_xlabel('Source Position', fontsize=11)
ax1.set_ylabel('Target Generation Step', fontsize=11)
ax1.set_title('Bahdanau (Additive) Alignment\nW_q·Q + W_k·K',
fontsize=12, fontweight='bold')
plt.colorbar(im1, ax=ax1, label='Attention Weight')
# 2. Luong Dot对齐
luong_dot = LuongAttention(hidden_dim, hidden_dim, score_mode='dot')
# 调整encoder维度以匹配（dot要求相同维度）
encoder_outputs_dot = torch.randn(batch_size, src_len, hidden_dim)
alignment_dot = []
with torch.no_grad():
for t in range(tgt_len):
query = decoder_hidden + 0.1 * torch.sin(target_positions[:, t:t+1] * 0.5).expand(-1, hidden_dim)
context, weights, scores = luong_dot(query, encoder_outputs_dot, encoder_outputs_dot)
alignment_dot.append(weights[0].numpy())
ax2 = axes[0, 1]
im2 = ax2.imshow(np.array(alignment_dot), cmap='plasma', aspect='auto', interpolation='nearest')
ax2.set_xlabel('Source Position', fontsize=11)
ax2.set_ylabel('Target Generation Step', fontsize=11)
ax2.set_title('Luong (Dot) Alignment\nQ^T·K',
fontsize=12, fontweight='bold')
plt.colorbar(im2, ax=ax2, label='Attention Weight')
# 3. Luong General对齐（可学习矩阵）
luong_general = LuongAttention(hidden_dim, hidden_dim, score_mode='general')
alignment_general = []
with torch.no_grad():
for t in range(tgt_len):
query = decoder_hidden + 0.1 * torch.sin(target_positions[:, t:t+1] * 0.5).expand(-1, hidden_dim)
context, weights, scores = luong_general(query, encoder_outputs_dot, encoder_outputs_dot)
alignment_general.append(weights[0].numpy())
ax3 = axes[0, 2]
im3 = ax3.imshow(np.array(alignment_general), cmap='coolwarm', aspect='auto', interpolation='nearest')
ax3.set_xlabel('Source Position', fontsize=11)
ax3.set_ylabel('Target Generation Step', fontsize=11)
ax3.set_title('Luong (General) Alignment\nQ^T·W·K',
fontsize=12, fontweight='bold')
plt.colorbar(im3, ax=ax3, label='Attention Weight')
# 4. K-Q-V寻址视角可视化
kqv_model = KeyQueryValuePerspective(dim=64, num_heads=1)
# 创建结构化数据：检索任务
seq_len = 10
x = torch.zeros(1, seq_len, 64)
# 设置特定位置存储特定信息
memory_slots = [2, 5, 8]
for idx, slot in enumerate(memory_slots):
x[0, slot, idx*20:(idx+1)*20] = 1.0 # 稀疏编码
with torch.no_grad():
output, attn_weights, (K, Q, V) = kqv_model(x)
ax4 = axes[1, 0]
# 可视化注意力矩阵（寻址矩阵）
attn_map = attn_weights[0].numpy()
im4 = ax4.imshow(attn_map, cmap='cividis', aspect='auto')
ax4.set_xlabel('Key Position (Memory Address)', fontsize=11)
ax4.set_ylabel('Query Position (Retrieval Request)', fontsize=11)
ax4.set_title('Key-Query-Value Addressing\nAttention as Memory Retrieval',
fontsize=12, fontweight='bold')
plt.colorbar(im4, ax=ax4, label='Retrieval Weight')
# 标记关键位置
for slot in memory_slots:
ax4.axvline(x=slot, color='red', linestyle='--', alpha=0.3)
ax4.axhline(y=slot, color='red', linestyle='--', alpha=0.3)
# 5. 评分函数对比（数值分布）
ax5 = axes[1, 1]
# 生成随机查询和键，计算不同评分函数的分布
sample_queries = torch.randn(1000, hidden_dim)
sample_keys = torch.randn(1000, hidden_dim)
# Dot product scores
dot_scores = (sample_queries * sample_keys).sum(dim=-1).numpy()
# Additive scores（模拟）
additive_layer = nn.Linear(hidden_dim * 2, 1)
combined = torch.cat([sample_queries, sample_keys], dim=-1)
additive_scores = additive_layer(combined).squeeze().detach().numpy()
ax5.hist(dot_scores, bins=50, alpha=0.6, label='Dot Product', color='steelblue', density=True)
ax5.hist(additive_scores, bins=50, alpha=0.6, label='Additive (Bahdanau)', color='darkorange', density=True)
ax5.set_xlabel('Score Value', fontsize=11)
ax5.set_ylabel('Density', fontsize=11)
ax5.set_title('Score Distribution Comparison', fontsize=12, fontweight='bold')
ax5.legend()
ax5.grid(True, alpha=0.3)
# 6. 计算效率对比（FLOPs估计）
ax6 = axes[1, 2]
dims = [64, 128, 256, 512, 1024]
seq_lens = [10, 50, 100, 200]
# 计算复杂度：Bahdanau O(d^2), Luong O(d) per step
bahdanau_flops = [d * d * sl for d in dims for sl in seq_lens]
luong_flops = [d * sl for d in dims for sl in seq_lens]
# 重塑为矩阵用于热力图
bahdanau_matrix = np.array(bahdanau_flops).reshape(len(dims), len(seq_lens))
luong_matrix = np.array(luong_flops).reshape(len(dims), len(seq_lens))
# 绘制效率比值（Bahdanau / Luong）
efficiency_ratio = bahdanau_matrix / luong_matrix
im6 = ax6.imshow(efficiency_ratio, cmap='YlOrRd', aspect='auto')
ax6.set_xticks(range(len(seq_lens)))
ax6.set_xticklabels(seq_lens)
ax6.set_yticks(range(len(dims)))
ax6.set_yticklabels(dims)
ax6.set_xlabel('Sequence Length', fontsize=11)
ax6.set_ylabel('Hidden Dimension', fontsize=11)
ax6.set_title('Computational Cost Ratio\nBahdanau / Luong (Higher=Slower)',
fontsize=12, fontweight='bold')
plt.colorbar(im6, ax=ax6, label='FLOP Ratio')
plt.tight_layout()
plt.savefig('attention_comparison.png', dpi=150, bbox_inches='tight')
print("Visualization saved to attention_comparison.png")
plt.show()
if __name__ == "__main__":
visualize_attention_mechanisms()10.2.2 多头自注意力的多头分割与输出拼接，缩放点积的Softmax温度缩放与注意力权重稀疏性分析
多头自注意力通过将查询、键、值向量在特征维度上均分为多个头，在每个子空间内独立执行缩放点积注意力，最终将各头输出拼接并线性投影。这种分割机制允许模型在不同表示子空间中同时关注来自不同位置的信息，形成对输入的多视角理解。每个头可以关注不同的语义或句法关系，如指代消解、句法依赖或语义角色标注。
缩放点积机制通过将点积结果除以特征维度的平方根进行温度缩放，防止在高维空间中点积值过大导致Softmax函数进入饱和区。这种缩放调节了注意力分布的熵，较大的缩放系数使分布更尖锐，聚焦于最相关的键；较小的系数使分布更平滑，保留更多候选键的信息。
注意力权重的稀疏性分析揭示，尽管Softmax输出在数学上稠密，实际训练后的注意力矩阵往往呈现近似稀疏结构。特定头仅强烈激活少数几个关键位置，其余权重接近零。这种涌现的稀疏性源于任务的结构性约束和梯度优化的动态特性，为后续的高效稀疏注意力算法提供了经验基础。
脚本：multihead_attention_sparse.py
Python
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脚本内容：多头自注意力实现与稀疏性分析
使用方式：直接运行 python multihead_attention_sparse.py
依赖包：torch, numpy, matplotlib, seaborn
功能说明：
1. 实现标准Multi-Head Attention（分割、并行计算、拼接）
2. 分析缩放点积中的温度缩放效应（Temperature Scaling）
3. 可视化注意力权重的稀疏性模式（头间差异、层间演化）
4. 展示多头如何捕捉不同子空间的关系
"""
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from typing import Optional, Tuple
import seaborn as sns
class MultiHeadAttention(nn.Module):
"""完整的多头自注意力实现，支持稀疏性分析"""
def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
super().__init__()
assert d_model % num_heads == 0, "d_model必须能被num_heads整除"
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_heads # 每头的维度
# 线性投影层：将输入投影到Q、K、V
self.W_q = nn.Linear(d_model, d_model, bias=False)
self.W_k = nn.Linear(d_model, d_model, bias=False)
self.W_v = nn.Linear(d_model, d_model, bias=False)
self.W_o = nn.Linear(d_model, d_model, bias=False) # 输出投影
self.dropout = nn.Dropout(dropout)
self.scale = np.sqrt(self.d_k) # 温度缩放因子
# 存储注意力权重用于可视化
self.attention_weights = None
def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None,
return_attention: bool = False):
"""
Args:
x: [batch, seq_len, d_model]
Returns:
output: [batch, seq_len, d_model]
attention: [batch, num_heads, seq_len, seq_len] (optional)
"""
batch_size, seq_len, _ = x.size()
# 1. 线性投影并分割为多头
# Q: [batch, seq, d_model] -> [batch, seq, num_heads, d_k] -> [batch, num_heads, seq, d_k]
Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
# 2. 缩放点积注意力
# scores: [batch, num_heads, seq, seq]
scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
# 应用掩码（因果掩码或填充掩码）
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
# Softmax获取注意力分布
attn_weights = torch.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)
if return_attention:
self.attention_weights = attn_weights.detach()
# 3. 加权聚合Value
# context: [batch, num_heads, seq, d_k]
context = torch.matmul(attn_weights, V)
# 4. 多头拼接与输出投影
# 交换维度：[batch, num_heads, seq, d_k] -> [batch, seq, num_heads, d_k]
context = context.transpose(1, 2).contiguous()
# 展平多头：[batch, seq, d_model]
context = context.view(batch_size, seq_len, self.d_model)
output = self.W_o(context)
return output, attn_weights if return_attention else output
class TemperatureScalingAnalysis:
"""分析温度缩放对注意力分布的影响"""
def __init__(self, d_k: int = 64):
self.d_k = d_k
def compute_attention_with_temperature(self, Q: torch.Tensor, K: torch.Tensor,
temperature_multiplier: float):
"""
使用不同的温度缩放因子计算注意力
temperature_multiplier: 相对于标准sqrt(d_k)的乘数
"""
effective_scale = np.sqrt(self.d_k) * temperature_multiplier
scores = torch.matmul(Q, K.transpose(-2, -1)) / effective_scale
attn = torch.softmax(scores, dim=-1)
return attn, scores
def analyze_entropy(self, attn_weights: torch.Tensor) -> float:
"""计算注意力分布的熵（稀疏性指标）"""
# 熵越低表示越稀疏（集中于少数位置）
entropy = -(attn_weights * torch.log(attn_weights + 1e-9)).sum(dim=-1).mean().item()
return entropy
def analyze_sparsity(self, attn_weights: torch.Tensor, threshold: float = 0.1) -> float:
"""计算注意力权重的稀疏度（低于threshold的比例）"""
sparse_ratio = (attn_weights < threshold).float().mean().item()
return sparse_ratio
class HeadSpecializationAnalyzer:
"""分析不同头的专业化模式"""
def __init__(self, num_heads: int):
self.num_heads = num_heads
self.head_patterns = {}
def analyze_head_behavior(self, attention_weights: torch.Tensor,
seq_tokens: Optional[list] = None):
"""
分析每个头的注意力模式
attention_weights: [batch, num_heads, seq, seq]
"""
batch_size, num_heads, seq_len, _ = attention_weights.size()
patterns = {}
for head_idx in range(num_heads):
head_attn = attention_weights[:, head_idx, :, :].mean(0).numpy() # 平均batch
# 检测对角线关注（位置依赖）
diagonal_focus = np.mean([head_attn[i, i] for i in range(min(20, seq_len))])
# 检测全局平均（均匀分布）
uniformity = np.std(head_attn.flatten())
# 检测局部窗口关注
window_focus = np.mean([head_attn[i, max(0,i-2):min(seq_len,i+3)].mean()
for i in range(seq_len)])
patterns[f'head_{head_idx}'] = {
'diagonal': diagonal_focus,
'uniformity': uniformity,
'local_window': window_focus,
'entropy': self._compute_entropy_numpy(head_attn)
}
return patterns
def _compute_entropy_numpy(self, attn_matrix: np.ndarray) -> float:
row_entropies = -(attn_matrix * np.log(attn_matrix + 1e-9)).sum(axis=-1)
return np.mean(row_entropies)
def visualize_multihead_sparse():
"""可视化多头注意力的稀疏性特性"""
# 设置参数
batch_size = 8
seq_len = 64
d_model = 256
num_heads = 8
# 创建模拟输入：合成结构化序列（包含重复模式、局部依赖和长期依赖）
x = torch.randn(batch_size, seq_len, d_model)
# 添加位置相关的结构（模拟真实语言的局部性和长程依赖）
for b in range(batch_size):
for i in range(seq_len):
# 局部平滑性
if i > 0:
x[b, i, :d_model//4] = x[b, i-1, :d_model//4] + 0.1 * torch.randn(d_model//4)
# 周期性长程依赖
if i >= 10:
x[b, i, d_model//4:d_model//2] = x[b, i-10, d_model//4:d_model//2] * 0.5
fig, axes = plt.subplots(2, 3, figsize=(16, 10))
# 1. 多头注意力权重热力图（展示不同头的专业化）
mha = MultiHeadAttention(d_model, num_heads)
mha.eval()
with torch.no_grad():
_, attn_weights = mha(x, return_attention=True)
# attn_weights: [batch, num_heads, seq, seq]
# 选择第一个样本，展示所有头的模式
sample_attn = attn_weights[0].numpy() # [num_heads, seq, seq]
# 绘制子图：选择4个代表性头
ax1 = axes[0, 0]
# 创建组合图：4个头的注意力矩阵
fig_attn, axes_attn = plt.subplots(2, 2, figsize=(10, 10))
for idx, ax in enumerate(axes_attn.flat):
im = ax.imshow(sample_attn[idx], cmap='viridis', aspect='auto', interpolation='nearest')
ax.set_title(f'Head {idx} Attention Pattern', fontsize=10)
ax.set_xlabel('Key Position')
ax.set_ylabel('Query Position')
plt.colorbar(im, ax=ax)
plt.tight_layout()
plt.savefig('head_specialization.png', dpi=150, bbox_inches='tight')
print("Head specialization saved to head_specialization.png")
# 在主图中显示第一个头的详细视图
im1 = ax1.imshow(sample_attn[0], cmap='plasma', aspect='auto', interpolation='nearest')
ax1.set_title('Head 0: Local/Diagonal Focus', fontsize=12, fontweight='bold')
ax1.set_xlabel('Key Position', fontsize=10)
ax1.set_ylabel('Query Position', fontsize=10)
plt.colorbar(im1, ax=ax1, label='Attention Weight')
# 2. 不同温度缩放的Softmax分布对比
ax2 = axes[0, 1]
temp_analyzer = TemperatureScalingAnalysis(d_k=d_model//num_heads)
# 生成随机Q、K
Q_demo = torch.randn(1, 1, 100, d_model//num_heads)
K_demo = torch.randn(1, 1, 100, d_model//num_heads)
temps = [0.5, 1.0, 2.0, 4.0] # 温度乘数（<1更尖锐，>1更平滑）
colors = plt.cm.plasma(np.linspace(0, 0.8, len(temps)))
for temp, color in zip(temps, colors):
attn, scores = temp_analyzer.compute_attention_with_temperature(
Q_demo, K_demo, temp
)
# 绘制注意力分布（第一个查询位置）
distribution = attn[0, 0, 0, :].numpy()
ax2.plot(distribution, label=f'Temp Multiplier={temp}',
color=color, linewidth=2, alpha=0.8)
ax2.set_xlabel('Key Position', fontsize=11)
ax2.set_ylabel('Attention Weight', fontsize=11)
ax2.set_title('Temperature Scaling Effect\non Attention Distribution',
fontsize=12, fontweight='bold')
ax2.legend()
ax2.grid(True, alpha=0.3)
# 3. 温度与稀疏性/熵的关系
ax3 = axes[0, 2]
temp_range = np.linspace(0.1, 5.0, 50)
entropies = []
max_weights = [] # 最大权重（稀疏性指标）
for t in temp_range:
attn, _ = temp_analyzer.compute_attention_with_temperature(Q_demo, K_demo, t)
entropy = temp_analyzer.analyze_entropy(attn)
max_w = attn.max(dim=-1)[0].mean().item()
entropies.append(entropy)
max_weights.append(max_w)
ax3_twin = ax3.twinx()
line1 = ax3.plot(temp_range, entropies, 'b-', linewidth=2, label='Entropy (diversity)')
line2 = ax3_twin.plot(temp_range, max_weights, 'r--', linewidth=2,
label='Max Weight (sparsity)')
ax3.set_xlabel('Temperature Multiplier', fontsize=11)
ax3.set_ylabel('Entropy (Higher=Uniform)', fontsize=11, color='b')
ax3_twin.set_ylabel('Max Attention Weight', fontsize=11, color='r')
ax3.set_title('Temperature vs Sparsity Trade-off', fontsize=12, fontweight='bold')
ax3.grid(True, alpha=0.3)
# 合并图例
lines = line1 + line2
labels = [l.get_label() for l in lines]
ax3.legend(lines, labels, loc='center right')
# 4. 多头稀疏性分析（头间差异）
ax4 = axes[1, 0]
analyzer = HeadSpecializationAnalyzer(num_heads)
patterns = analyzer.analyze_head_behavior(attn_weights)
head_indices = list(range(num_heads))
diagonal_scores = [patterns[f'head_{i}']['diagonal'] for i in head_indices]
entropies_heads = [patterns[f'head_{i}']['entropy'] for i in head_indices]
x_pos = np.arange(num_heads)
width = 0.35
bars1 = ax4.bar(x_pos - width/2, diagonal_scores, width,
label='Diagonal Focus', color='steelblue', alpha=0.8)
ax4_twin = ax4.twinx()
bars2 = ax4_twin.bar(x_pos + width/2, entropies_heads, width,
label='Entropy', color='darkorange', alpha=0.8)
ax4.set_xlabel('Head Index', fontsize=11)
ax4.set_ylabel('Diagonal Focus Score', fontsize=11, color='steelblue')
ax4_twin.set_ylabel('Attention Entropy', fontsize=11, color='darkorange')
ax4.set_title('Head Specialization Patterns', fontsize=12, fontweight='bold')
ax4.set_xticks(x_pos)
ax4.grid(True, alpha=0.3, axis='y')
# 添加文本注释识别专业化头
for i, (diag, ent) in enumerate(zip(diagonal_scores, entropies_heads)):
if diag > 0.3:
ax4.annotate('Local', (i, diag), textcoords="offset points",
xytext=(0,5), ha='center', fontsize=8, color='blue')
if ent < 2.0:
ax4.annotate('Sparse', (i, ent), textcoords="offset points",
xytext=(0,5), ha='center', fontsize=8, color='red')
# 5. 注意力权重分布（稀疏性可视化）
ax5 = axes[1, 1]
# 收集所有头的权重分布
all_weights = attn_weights.flatten().numpy()
# 绘制分布直方图
ax5.hist(all_weights, bins=100, color='purple', alpha=0.7, density=True,
label='Attention Weights')
ax5.axvline(x=0.01, color='red', linestyle='--', label='Sparsity Threshold (0.01)')
ax5.axvline(x=all_weights.mean(), color='green', linestyle='-',
label=f'Mean ({all_weights.mean():.4f})')
# 计算稀疏度
sparsity = (all_weights < 0.01).mean()
ax5.text(0.02, ax5.get_ylim()[1]*0.8, f'Sparsity: {sparsity:.2%}\n(weights < 0.01)',
bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
ax5.set_xlabel('Attention Weight Value', fontsize=11)
ax5.set_ylabel('Density', fontsize=11)
ax5.set_title('Attention Weight Distribution\n(Emergent Sparsity)',
fontsize=12, fontweight='bold')
ax5.set_yscale('log')
ax5.legend()
ax5.grid(True, alpha=0.3)
# 6. 层间稀疏性演化（模拟多层Transformer）
ax6 = axes[1, 2]
num_layers = 6
layer_entropies = []
layer_sparsities = []
# 模拟不同层的稀疏性模式（通常深层更稀疏）
for layer in range(num_layers):
# 深层倾向于更稀疏（更低的熵）
base_entropy = 3.5 - layer * 0.4 # 递减
noise = np.random.randn(batch_size, num_heads, seq_len, seq_len) * 0.5
sim_attn = torch.softmax(torch.randn(batch_size, num_heads, seq_len, seq_len)
* (1.0 / (1 + layer*0.2)), dim=-1)
ent = -(sim_attn * torch.log(sim_attn + 1e-9)).sum(dim=-1).mean().item()
spar = (sim_attn < 0.05).float().mean().item()
layer_entropies.append(ent)
layer_sparsities.append(spar)
ax6.plot(range(num_layers), layer_entropies, 'o-', label='Entropy',
color='navy', linewidth=2, markersize=8)
ax6_twin = ax6.twinx()
ax6_twin.plot(range(num_layers), [s*100 for s in layer_sparsities], 's--',
label='Sparsity %', color='crimson', linewidth=2, markersize=8)
ax6.set_xlabel('Layer Depth', fontsize=11)
ax6.set_ylabel('Average Entropy', fontsize=11, color='navy')
ax6_twin.set_ylabel('Sparsity %', fontsize=11, color='crimson')
ax6.set_title('Sparsity Evolution Across Layers', fontsize=12, fontweight='bold')
ax6.set_xticks(range(num_layers))
ax6.grid(True, alpha=0.3)
ax6.legend(loc='upper left')
ax6_twin.legend(loc='upper right')
plt.tight_layout()
plt.savefig('multihead_sparse_analysis.png', dpi=150, bbox_inches='tight')
print("Main visualization saved to multihead_sparse_analysis.png")
plt.show()
if __name__ == "__main__":
visualize_multihead_sparse()10.2.3 位置编码：正弦/余弦绝对位置编码的波长设计，可学习位置编码与RoPE旋转位置编码的复数域实现
位置编码为Transformer引入序列顺序信息，弥补自注意力本身的位置无关性。正弦余弦绝对位置编码通过不同频率的正弦和余弦函数为每个位置生成唯一编码，波长设计遵循几何级数分布，从高频短波长到低频长波长覆盖多尺度位置关系。这种设计使模型能够通过线性组合学习相对位置信息，并通过正弦函数的周期性处理超出训练长度的序列。
可学习位置编码将位置向量作为模型参数随机初始化，通过反向传播学习最优位置表示。这种方法提供最大灵活性，能够适应特定任务的特定位置模式，但需要更多训练数据，且在推断时难以外推至未见过的更长序列。
旋转位置编码通过将查询和键向量表示为复数，利用复数乘法实现旋转操作，将相对位置信息编码为旋转相位。在复数域中，位置编码转化为向量的旋转变换，使点积注意力自然体现相对位置依赖性。这种实现保持了Transformer的数学一致性，同时在长序列外推和相对位置感知方面表现优异，已成为现代大语言模型的标准配置。
脚本：positional_encoding_comparison.py
Python
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脚本内容：位置编码机制对比与实现
使用方式：直接运行 python positional_encoding_comparison.py
依赖包：torch, numpy, matplotlib
功能说明：
1. 实现Sinusoidal、Learnable、RoPE三种位置编码
2. 可视化正弦余弦编码的多尺度波长设计
3. 展示RoPE的复数域旋转实现与外推能力
4. 对比不同编码在长程依赖和长度外推上的表现
"""
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from typing import Optional
class SinusoidalPositionalEncoding(nn.Module):
"""正弦余弦绝对位置编码（Vaswani et al. 原版）"""
def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
# 计算位置编码矩阵
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
# 波长几何级数设计
div_term = torch.exp(torch.arange(0, d_model, 2).float() *
(-np.log(10000.0) / d_model))
# 正弦用于偶数维度，余弦用于奇数维度
pe[:, 0::2] = torch.sin(position * div_term) # 2i
pe[:, 1::2] = torch.cos(position * div_term) # 2i+1
pe = pe.unsqueeze(0).transpose(0, 1) # [max_len, 1, d_model]
self.register_buffer('pe', pe)
# 存储波长用于可视化
self.wavelengths = 2 * np.pi / div_term.numpy()
def forward(self, x: torch.Tensor):
x = x + self.pe[:x.size(0), :]
return self.dropout(x)
class LearnablePositionalEncoding(nn.Module):
"""可学习位置编码（BERT风格）"""
def __init__(self, d_model: int, max_len: int = 512, dropout: float = 0.1):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
# 将位置嵌入作为可学习参数
self.pe = nn.Embedding(max_len, d_model)
self.pe.weight.data.normal_(mean=0.0, std=0.02) # 小随机初始化
self.max_len = max_len
def forward(self, x: torch.Tensor, positions: Optional[torch.Tensor] = None):
seq_len = x.size(0)
if positions is None:
positions = torch.arange(seq_len, device=x.device).unsqueeze(1).expand(-1, x.size(1))
pos_embed = self.pe(positions)
x = x + pos_embed
return self.dropout(x)
class RotaryPositionalEmbedding(nn.Module):
"""旋转位置编码RoPE的复数域实现"""
def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0):
super().__init__()
self.dim = dim
# 计算旋转角度（频率）
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
self.register_buffer('inv_freq', inv_freq)
# 预计算旋转矩阵（用于效率）
self.max_seq_len = max_seq_len
self._precompute_cache()
def _precompute_cache(self):
"""预计算余弦和正弦值（旋转矩阵元素）"""
t = torch.arange(self.max_seq_len, device=self.inv_freq.device)
freqs = torch.einsum('i,j->ij', t, self.inv_freq) # [seq_len, dim//2]
# 在复数域中，旋转矩阵表示为cos(θ) + i*sin(θ)
self.register_buffer('cos_cached', freqs.cos())
self.register_buffer('sin_cached', freqs.sin())
def apply_rotary_pos_emb(self, x: torch.Tensor, seq_len: int):
"""
将旋转位置编码应用于输入（复数域旋转实现）
x: [batch, num_heads, seq_len, head_dim]
"""
# 将x视为复数对（实部和虚部）
# x = [x0, x1, x2, x3, ...] 转换为 [(x0,x1), (x2,x3), ...]
x1 = x[..., ::2] # 偶数索引（实部）
x2 = x[..., 1::2] # 奇数索引（虚部）
# 获取对应位置的旋转矩阵元素
cos = self.cos_cached[:seq_len].unsqueeze(0).unsqueeze(0) # [1, 1, seq, dim//2]
sin = self.sin_cached[:seq_len].unsqueeze(0).unsqueeze(0)
# 复数乘法：(a+ib)(cos+isin) = (a*cos - b*sin) + i(a*sin + b*cos)
y1 = x1 * cos - x2 * sin # 新的实部
y2 = x1 * sin + x2 * cos # 新的虚部
# 交错合并回原始格式
y = torch.stack([y1, y2], dim=-1).flatten(-2)
return y
def forward(self, q: torch.Tensor, k: torch.Tensor):
"""对查询和键应用RoPE"""
seq_len = q.size(-2)
q_rotated = self.apply_rotary_pos_emb(q, seq_len)
k_rotated = self.apply_rotary_pos_emb(k, seq_len)
return q_rotated, k_rotated
def visualize_positional_encodings():
"""可视化不同位置编码的特性和行为"""
d_model = 128
max_len_viz = 256
seq_len = 100
fig, axes = plt.subplots(2, 3, figsize=(16, 10))
# 1. 正弦编码的波长设计（多尺度分析）
ax1 = axes[0, 0]
sin_pe = SinusoidalPositionalEncoding(d_model, max_len=max_len_viz)
# 展示不同维度的波长
dims_to_plot = [0, 8, 16, 32, 64] # 不同频率的维度
colors = plt.cm.plasma(np.linspace(0, 0.8, len(dims_to_plot)))
for dim_idx, color in zip(dims_to_plot, colors):
wavelength = sin_pe.wavelengths[dim_idx // 2]
positions = np.arange(max_len_viz)
values = sin_pe.pe[:, 0, dim_idx].numpy()
ax1.plot(positions, values, label=f'Dim {dim_idx} (λ={wavelength:.1f})',
color=color, linewidth=2, alpha=0.8)
ax1.set_xlabel('Position', fontsize=11)
ax1.set_ylabel('Encoding Value', fontsize=11)
ax1.set_title('Sinusoidal Encoding: Multi-Scale Wavelength Design\n'
'Geometric Progression from Short to Long',
fontsize=12, fontweight='bold')
ax1.legend(loc='upper right', fontsize=8)
ax1.grid(True, alpha=0.3)
# 2. 正弦编码的周期性和线性变换特性
ax2 = axes[0, 1]
pe_matrix = sin_pe.pe[:128, 0, :].numpy() # [128, d_model]
im = ax2.imshow(pe_matrix, aspect='auto', cmap='RdBu_r', vmin=-1, vmax=1)
ax2.set_xlabel('Dimension', fontsize=11)
ax2.set_ylabel('Position', fontsize=11)
ax2.set_title('Sinusoidal Encoding Matrix\n'
'Positional Information as Phase Shifts',
fontsize=12, fontweight='bold')
plt.colorbar(im, ax=ax2, label='Value')
# 添加波长注释
for dim in [0, 32, 64, 96]:
wl = sin_pe.wavelengths[dim // 2] if dim % 2 == 0 else sin_pe.wavelengths[dim // 2]
if wl < 128:
ax2.axvline(x=dim, color='yellow', linestyle='--', alpha=0.5)
ax2.text(dim+2, 120, f'λ={wl:.0f}', color='yellow', fontsize=8, rotation=90)
# 3. 可学习编码 vs 正弦编码的分布对比
ax3 = axes[0, 2]
learn_pe = LearnablePositionalEncoding(d_model, max_len=max_len_viz)
# 获取可学习嵌入
learned_matrix = learn_pe.pe.weight.detach().numpy()[:128, :]
sinusoidal_matrix = pe_matrix[:128, :128] # 取相同维度
# 计算每对位置的余弦相似度（展示位置关系学习）
sim_learned = np.dot(learned_matrix, learned_matrix.T)
sim_sinusoidal = np.dot(sinusoidal_matrix, sinusoidal_matrix.T)
# 绘制相似度矩阵的对角线（相邻位置关系）
ax3.plot(np.diag(sim_learned, k=1), label='Learned (Adjacent Similarity)',
color='crimson', linewidth=2, marker='o', markersize=4)
ax3.plot(np.diag(sim_sinusoidal, k=1), label='Sinusoidal (Theoretical)',
color='navy', linewidth=2, linestyle='--', marker='s', markersize=4)
ax3.set_xlabel('Position', fontsize=11)
ax3.set_ylabel('Cosine Similarity with Next Position', fontsize=11)
ax3.set_title('Position Relationship: Learned vs Theoretical\n'
'Learned embeddings adapt to data statistics',
fontsize=12, fontweight='bold')
ax3.legend()
ax3.grid(True, alpha=0.3)
# 4. RoPE的旋转机制可视化（复数域）
ax4 = axes[1, 0]
rope = RotaryPositionalEmbedding(d_model)
# 可视化旋转角度随位置和维度的变化
freqs = rope.inv_freq.numpy()
positions = np.arange(0, 100, 10)
# 选择几个维度展示旋转角度累积
selected_dims = [0, 4, 8, 16]
for dim_idx in selected_dims:
angles = positions * freqs[dim_idx] # θ = pos * frequency
ax4.plot(positions, angles, 'o-', label=f'Dim Pair {dim_idx}',
linewidth=2, markersize=6)
ax4.set_xlabel('Position', fontsize=11)
ax4.set_ylabel('Rotation Angle (radians)', fontsize=11)
ax4.set_title('RoPE: Rotation Angle Accumulation\n'
'Each dimension pair rotates at different speed',
fontsize=12, fontweight='bold')
ax4.legend()
ax4.grid(True, alpha=0.3)
# 添加2π周期标记
for multiple in [1, 2, 3]:
ax4.axhline(y=multiple * 2 * np.pi, color='red', linestyle=':', alpha=0.5)
ax4.text(90, multiple * 2 * np.pi + 0.2, f'{multiple}·2π', fontsize=9, color='red')
# 5. RoPE的相对位置编码特性（注意力矩阵）
ax5 = axes[1, 1]
# 模拟查询和键向量
num_positions = 64
head_dim = 32
# 创建简单模式：每个位置有独特但相关的向量
q_base = torch.randn(1, 1, num_positions, head_dim)
k_base = q_base.clone() # 相同内容，测试位置敏感性
# 应用RoPE
rope_local = RotaryPositionalEmbedding(head_dim, max_seq_len=num_positions)
q_rot, k_rot = rope_local(q_base, k_base)
# 计算注意力分数（体现相对位置）
scores = torch.matmul(q_rot, k_rot.transpose(-2, -1)) / np.sqrt(head_dim)
scores_np = scores[0, 0].numpy()
im5 = ax5.imshow(scores_np, cmap='coolwarm', aspect='auto')
ax5.set_xlabel('Key Position', fontsize=11)
ax5.set_ylabel('Query Position', fontsize=11)
ax5.set_title('RoPE Attention Scores\n'
'Diagonal dominance = Relative Position Bias',
fontsize=12, fontweight='bold')
plt.colorbar(im5, ax=ax5, label='Attention Score')
# 添加对角线注释（相对位置）
for offset in [-20, 0, 20]:
if 0 <= num_positions//2 + offset < num_positions:
ax5.annotate(f'Δ={offset}',
xy=(num_positions//2+offset, num_positions//2),
xytext=(num_positions//2+offset, num_positions//2+10),
arrowprops=dict(arrowstyle='->', color='yellow'),
color='yellow', fontsize=9, ha='center')
# 6. 长度外推能力对比（关键特性）
ax6 = axes[1, 2]
# 训练长度
train_len = 100
# 测试外推到更长长度
test_len = 200
# 正弦编码的外推（理论上无限外推）
sin_full = SinusoidalPositionalEncoding(d_model, max_len=test_len)
sin_train = sin_full.pe[:train_len, 0, :].numpy()
sin_test = sin_full.pe[train_len:, 0, :].numpy()
# 计算训练位置与测试位置的相似度（外推质量）
similarity_matrix = np.dot(sin_train, sin_test.T) # [train_len, test_len-train_len]
# 可学习编码无法直接外推（只能评估训练内的）
learned_train = learn_pe.pe.weight[:train_len, :].detach().numpy()
learned_test_sim = np.dot(learned_train, learned_train[-10:].T) # 自我相似
# RoPE外推（旋转连续性）
rope_train_angles = np.arange(train_len)[:, None] * rope.inv_freq.numpy()[None, :]
rope_test_angles = np.arange(train_len, test_len)[:, None] * rope.inv_freq.numpy()[None, :]
# 计算角度差异的连续性
angle_diff = np.abs(rope_test_angles[0] - rope_train_angles[-1])
# 可视化外推连续性
x_range = np.arange(train_len, test_len)
# 取第一个维度的角度连续性
continuity = np.abs(rope_test_angles[:, 0] - rope_train_angles[-1, 0])
ax6.plot(x_range, continuity, 'b-', linewidth=2, label='RoPE Angle Continuity')
ax6.axvline(x=train_len, color='red', linestyle='--', alpha=0.7, label='Training Boundary')
ax6.axhline(y=2*np.pi, color='orange', linestyle=':', alpha=0.5, label='2π Period')
# 绘制正弦编码的外推相似度（周期性重置）
periodicity = np.abs(np.sin(rope_test_angles[:, 0]) - np.sin(rope_train_angles[-1, 0]))
ax6.plot(x_range, periodicity * 2 * np.pi, 'g--', linewidth=2,
label='Sinusoidal Periodicity', alpha=0.7)
ax6.set_xlabel('Extrapolated Position', fontsize=11)
ax6.set_ylabel('Angular Continuity / Periodicity', fontsize=11)
ax6.set_title('Length Extrapolation Capability\n'
'RoPE maintains continuous rotation beyond training',
fontsize=12, fontweight='bold')
ax6.legend()
ax6.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('positional_encoding_comparison.png', dpi=150, bbox_inches='tight')
print("Visualization saved to positional_encoding_comparison.png")
plt.show()
def demonstrate_extrapolation_numerical():
"""数值演示RoPE的外推能力"""
print("\n=== RoPE Extrapolation Demonstration ===")
dim = 64
rope = RotaryPositionalEmbedding(dim, max_seq_len=1000)
# 训练时见过的位置
train_pos = 100
# 未见过的新位置
test_pos = 500
# 创建测试向量
x = torch.randn(1, 1, 1, dim)
# 应用RoPE到两个位置
angles_train = torch.tensor([[train_pos]]) * rope.inv_freq.unsqueeze(0)
angles_test = torch.tensor([[test_pos]]) * rope.inv_freq.unsqueeze(0)
# 计算旋转后的点积（体现相对位置）
# RoPE的关键特性：<R(pos_q)*q, R(pos_k)*k> 依赖于 (pos_q - pos_k)
rel_effect = test_pos - train_pos
print(f"Training position: {train_pos}")
print(f"Test position: {test_pos}")
print(f"Relative distance: {rel_effect}")
print(f"Rotation angle difference (first dim): {rel_effect * rope.inv_freq[0].item():.4f} rad")
print(f"Periodicity (2π intervals): {(rel_effect * rope.inv_freq[0].item()) / (2*np.pi):.2f}")
if __name__ == "__main__":
visualize_positional_encodings()
demonstrate_extrapolation_numerical()10.2.4 Transformer编码器-解码器架构：残差连接的Pre-LN与Post-LN布局对比，FFN的GELU激活与中间维度扩展比例
Transformer编码器-解码器架构通过堆叠自注意力层和交叉注意力层构建深度序列转换模型。残差连接在子层周围创建梯度高速公路，使误差信号能够绕过非线性变换直接传播。层归一化位置的布局选择对优化稳定性产生根本性影响：后归一化布局将归一化置于残差连接之后，原始子层输出直接参与求和，虽增强表达能力但导致深层梯度衰减；前归一化布局将归一化置于子层输入，稳定深层训练但可能限制模型容量。
前馈神经网络子层采用扩展-收缩结构，中间维度通常为输入维度的四倍，通过升维投影增加非线性变换的表达能力。高斯误差线性单元激活函数引入平滑的非线性门控机制，相比整流线性单元提供更连续的梯度流，在深层模型中表现出更优的优化特性。这种维度扩展与收缩模式模拟了核方法中的特征空间映射，将低维输入投影至高维空间执行非线性分离后再压缩还原。
脚本：transformer_architecture_prepost.py
Python
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脚本内容：Transformer架构变体（Pre-LN vs Post-LN）与FFN实现
使用方式：直接运行 python transformer_architecture_prepost.py
依赖包：torch, numpy, matplotlib
功能说明：
1. 实现完整的Transformer Encoder-Decoder（Pre-LN和Post-LN两种布局）
2. 对比残差连接与层归一化不同顺序的训练稳定性
3. 可视化梯度流与隐藏状态分布
4. 分析FFN中间层维度扩展（4d）与GELU激活的影响
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from typing import Optional, Tuple
class GELUActivation(nn.Module):
"""GELU激活函数的精确与近似实现"""
def __init__(self, approximate: str = 'none'):
super().__init__()
self.approximate = approximate
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.approximate == 'tanh':
# 近似版本（计算更快）
return 0.5 * x * (1.0 + torch.tanh(
np.sqrt(2.0 / np.pi) * (x + 0.044715 * torch.pow(x, 3))
))
else:
# 精确版本（基于误差函数）
return 0.5 * x * (1.0 + torch.erf(x / np.sqrt(2.0)))
class FeedForwardNetwork(nn.Module):
"""
FFN实现：扩展-收缩结构，中间维度4d
"""
def __init__(self, d_model: int, d_ff: Optional[int] = None,
dropout: float = 0.1, use_gelu: bool = True):
super().__init__()
self.d_model = d_model
self.d_ff = d_ff or 4 * d_model # 默认4倍扩展
# 扩展投影
self.w_1 = nn.Linear(d_model, self.d_ff)
# 收缩投影
self.w_2 = nn.Linear(self.d_ff, d_model)
self.activation = GELUActivation() if use_gelu else nn.ReLU()
self.dropout = nn.Dropout(dropout)
# 存储中间激活用于可视化
self.intermediate_activation = None
def forward(self, x: torch.Tensor) -> torch.Tensor:
# 升维
hidden = self.w_1(x) # [batch, seq, 4d]
hidden = self.activation(hidden)
self.intermediate_activation = hidden.detach()
hidden = self.dropout(hidden)
# 降维
output = self.w_2(hidden) # [batch, seq, d]
return output
class TransformerEncoderLayer(nn.Module):
"""支持Pre-LN和Post-LN的编码器层"""
def __init__(self, d_model: int, num_heads: int, d_ff: Optional[int] = None,
dropout: float = 0.1, norm_type: str = 'pre'):
super().__init__()
self.norm_type = norm_type # 'pre' 或 'post'
# 自注意力
self.self_attn = nn.MultiheadAttention(d_model, num_heads, dropout=dropout)
# FFN
self.ffn = FeedForwardNetwork(d_model, d_ff, dropout)
# 层归一化
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
# Dropout
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
# 存储统计信息用于可视化
self.attn_weights = None
self.grad_norms = {}
def forward(self, src: torch.Tensor, src_mask: Optional[torch.Tensor] = None,
is_training: bool = True):
"""
Args:
src: [seq_len, batch, d_model]
"""
# 自注意力块
if self.norm_type == 'pre':
# Pre-LN: Norm -> Sub-layer -> Residual
src2 = self.norm1(src)
src2, attn_weights = self.self_attn(src2, src2, src2, attn_mask=src_mask)
self.attn_weights = attn_weights
src = src + self.dropout1(src2) # 残差连接
# FFN块
src2 = self.norm2(src)
src2 = self.ffn(src2)
src = src + self.dropout2(src2)
else:
# Post-LN: Sub-layer -> Residual -> Norm
src2, attn_weights = self.self_attn(src, src, src, attn_mask=src_mask)
self.attn_weights = attn_weights
src = src + self.dropout1(src2)
src = self.norm1(src)
# FFN块
src2 = self.ffn(src)
src = src + self.dropout2(src2)
src = self.norm2(src)
return src
class TransformerEncoder(nn.Module):
"""堆叠的Transformer编码器"""
def __init__(self, num_layers: int, d_model: int, num_heads: int,
d_ff: Optional[int] = None, dropout: float = 0.1,
norm_type: str = 'pre'):
super().__init__()
self.norm_type = norm_type
self.layers = nn.ModuleList([
TransformerEncoderLayer(d_model, num_heads, d_ff, dropout, norm_type)
for _ in range(num_layers)
])
# 最终层归一化（Pre-LN需要）
if norm_type == 'pre':
self.final_norm = nn.LayerNorm(d_model)
else:
self.final_norm = nn.Identity()
self.num_layers = num_layers
def forward(self, src: torch.Tensor, src_mask: Optional[torch.Tensor] = None):
# 收集每层的统计信息
layer_outputs = []
layer_grad_norms = []
x = src
for idx, layer in enumerate(self.layers):
x = layer(x, src_mask)
layer_outputs.append(x)
# 计算梯度范数（用于分析稳定性）
if x.requires_grad and self.training:
def hook_fn(grad, layer_idx=idx):
layer_grad_norms.append((layer_idx, grad.norm().item()))
x.register_hook(hook_fn)
x = self.final_norm(x)
return x, layer_outputs
class StabilityAnalyzer:
"""分析Pre-LN与Post-LN的稳定性差异"""
def __init__(self, num_layers: int, d_model: int, num_heads: int):
self.num_layers = num_layers
self.d_model = d_model
self.results = {'pre': {}, 'post': {}}
def analyze_gradient_flow(self, model_type: str, num_iters: int = 100):
"""分析梯度流经各层的情况"""
encoder = TransformerEncoder(
self.num_layers, self.d_model, num_heads,
norm_type=model_type
)
encoder.train()
optimizer = torch.optim.Adam(encoder.parameters(), lr=0.001)
layer_grad_norms = [[] for _ in range(self.num_layers)]
for _ in range(num_iters):
# 模拟输入
src = torch.randn(20, 8, self.d_model, requires_grad=True)
optimizer.zero_grad()
output, _ = encoder(src)
# 从最后一层反向传播损失
loss = output.mean()
loss.backward()
# 收集各层梯度范数
for name, param in encoder.named_parameters():
if 'layers.' in name and param.grad is not None:
# 提取层索引
layer_idx = int(name.split('.')[1])
grad_norm = param.grad.norm().item()
layer_grad_norms[layer_idx].append(grad_norm)
optimizer.step()
# 计算统计信息
avg_grads = [np.mean(grads) if grads else 0 for grads in layer_grad_norms]
std_grads = [np.std(grads) if grads else 0 for grads in layer_grad_norms]
self.results[model_type]['grad_mean'] = avg_grads
self.results[model_type]['grad_std'] = std_grads
return avg_grads, std_grads
def analyze_hidden_dynamics(self, model_type: str):
"""分析隐藏状态的分布演化"""
encoder = TransformerEncoder(
self.num_layers, self.d_model, num_heads,
norm_type=model_type
)
encoder.eval()
with torch.no_grad():
src = torch.randn(50, 4, self.d_model)
_, layer_outputs = encoder(src)
# 计算每层输出的统计特征
means = []
stds = []
max_vals = []
for output in layer_outputs:
means.append(output.mean().item())
stds.append(output.std().item())
max_vals.append(output.abs().max().item())
self.results[model_type]['hidden_mean'] = means
self.results[model_type]['hidden_std'] = stds
self.results[model_type]['hidden_max'] = max_vals
return means, stds
def visualize_architecture_comparison():
"""可视化Pre-LN与Post-LN的详细对比"""
num_layers = 12
d_model = 256
num_heads = 8
fig, axes = plt.subplots(2, 3, figsize=(16, 10))
# 1. 架构示意图对比
ax1 = axes[0, 0]
ax1.set_xlim(0, 10)
ax1.set_ylim(0, 10)
ax1.axis('off')
ax1.set_title('Architecture Layout Comparison', fontsize=12, fontweight='bold')
# Pre-LN示意图
ax1.text(2.5, 9, 'Pre-LN', ha='center', fontsize=11, fontweight='bold')
pre_blocks = [
('Input', 8.5, 'lightblue'),
('Norm', 7.5, 'lightyellow'),
('Self-Attn', 6.5, 'lightgreen'),
('Residual (+)', 5.5, 'lightcoral'),
('Norm', 4.5, 'lightyellow'),
('FFN', 3.5, 'lightgreen'),
('Residual (+)', 2.5, 'lightcoral'),
('Output', 1.5, 'lightblue')
]
y_pos = 8.5
for label, y, color in pre_blocks:
ax1.add_patch(plt.Rectangle((0.5, y-0.3), 4, 0.6, facecolor=color, edgecolor='black'))
ax1.text(2.5, y, label, ha='center', va='center', fontsize=9)
if y > 1.5:
ax1.arrow(2.5, y-0.4, 0, -0.2, head_width=0.2, head_length=0.1, fc='black')
# Post-LN示意图
ax1.text(7.5, 9, 'Post-LN', ha='center', fontsize=11, fontweight='bold')
post_blocks = [
('Input', 8.5, 'lightblue'),
('Self-Attn', 7.5, 'lightgreen'),
('Residual (+)', 6.5, 'lightcoral'),
('Norm', 5.5, 'lightyellow'),
('FFN', 4.5, 'lightgreen'),
('Residual (+)', 3.5, 'lightcoral'),
('Norm', 2.5, 'lightyellow'),
('Output', 1.5, 'lightblue')
]
for label, y, color in post_blocks:
ax1.add_patch(plt.Rectangle((5.5, y-0.3), 4, 0.6, facecolor=color, edgecolor='black'))
ax1.text(7.5, y, label, ha='center', va='center', fontsize=9)
if y > 1.5:
ax1.arrow(7.5, y-0.4, 0, -0.2, head_width=0.2, head_length=0.1, fc='black')
# 2. 梯度流分析
ax2 = axes[0, 1]
analyzer = StabilityAnalyzer(num_layers, d_model, num_heads)
# 模拟梯度数据（基于文献观察）
pre_grads = [1.0 - i*0.02 for i in range(num_layers)] # Pre-LN梯度较均匀
post_grads = [0.3 + 0.7 * (i/num_layers)**2 for i in range(num_layers)] # Post-LN深层梯度大
x = np.arange(num_layers)
width = 0.35
bars1 = ax2.bar(x - width/2, pre_grads, width, label='Pre-LN',
color='steelblue', alpha=0.8)
bars2 = ax2.bar(x + width/2, post_grads, width, label='Post-LN',
color='coral', alpha=0.8)
ax2.set_xlabel('Layer Index', fontsize=11)
ax2.set_ylabel('Gradient Norm (Normalized)', fontsize=11)
ax2.set_title('Gradient Flow: Layer-wise Distribution\n'
'Pre-LN: Uniform | Post-LN: Increasing with Depth',
fontsize=12, fontweight='bold')
ax2.set_xticks(x[::2])
ax2.legend()
ax2.grid(True, alpha=0.3, axis='y')
# 3. 隐藏状态分布演化
ax3 = axes[0, 2]
# 模拟隐藏状态统计
pre_means = [0.0] * num_layers
pre_stds = [1.0 + i*0.1 for i in range(num_layers)] # Pre-LN方差增长
post_means = [0.0] * num_layers
post_stds = [1.0] * num_layers # Post-LN方差稳定
ax3.plot(x, pre_stds, 'o-', label='Pre-LN Std', color='steelblue', linewidth=2)
ax3.plot(x, post_stds, 's-', label='Post-LN Std', color='coral', linewidth=2)
ax3.axhline(y=1.0, color='black', linestyle='--', alpha=0.3, label='Ideal')
ax3.set_xlabel('Layer Depth', fontsize=11)
ax3.set_ylabel('Hidden State Std Dev', fontsize=11)
ax3.set_title('Hidden State Variance Growth\n'
'Pre-LN: Exploding Variance | Post-LN: Stable',
fontsize=12, fontweight='bold')
ax3.legend()
ax3.grid(True, alpha=0.3)
# 4. FFN中间层激活可视化（4d扩展）
ax4 = axes[1, 0]
ffn = FeedForwardNetwork(d_model, d_ff=4*d_model)
ffn.eval()
# 模拟不同输入下的中间层激活
test_input = torch.randn(100, d_model) # 100个样本
with torch.no_grad():
_ = ffn(test_input)
intermediate = ffn.intermediate_activation # [100, 4d]
# 可视化激活分布
sample_dims = np.random.choice(4*d_model, 50, replace=False)
activation_data = intermediate[:, sample_dims].numpy()
im4 = ax4.imshow(activation_data.T, aspect='auto', cmap='viridis',
vmin=-3, vmax=3, interpolation='nearest')
ax4.set_xlabel('Sample Index', fontsize=11)
ax4.set_ylabel('FFN Hidden Dimension (subset)', fontsize=11)
ax4.set_title('FFN Intermediate Activations (4× expansion)\n'
'Sparse activation pattern (mixture of experts effect)',
fontsize=12, fontweight='bold')
plt.colorbar(im4, ax=ax4, label='Activation Value')
# 统计稀疏性
sparsity = (intermediate == 0).float().mean().item()
ax4.text(5, 5, f'Sparsity: {sparsity:.2%}', color='white',
bbox=dict(boxstyle='round', facecolor='red', alpha=0.7),
fontsize=10, fontweight='bold')
# 5. GELU vs ReLU对比
ax5 = axes[1, 1]
x_range = np.linspace(-3, 3, 1000)
relu = np.maximum(0, x_range)
gelu = 0.5 * x_range * (1 + np.tanh(
np.sqrt(2/np.pi) * (x_range + 0.044715 * x_range**3)
))
ax5.plot(x_range, relu, label='ReLU', color='crimson',
linewidth=2.5, linestyle='--', alpha=0.8)
ax5.plot(x_range, gelu, label='GELU', color='forestgreen',
linewidth=2.5, alpha=0.8)
ax5.plot(x_range, x_range * (1 - np.exp(-x_range**2)),
label='Swish (SiLU)', color='blue',
linewidth=2, linestyle=':', alpha=0.6)
# 标记关键区域
ax5.axvline(x=0, color='black', linestyle=':', alpha=0.3)
ax5.fill_between(x_range, relu, gelu, where=(x_range < 0),
alpha=0.3, color='yellow', label='Negative region difference')
ax5.set_xlabel('Input Value', fontsize=11)
ax5.set_ylabel('Activation Output', fontsize=11)
ax5.set_title('Activation Functions: GELU vs ReLU\n'
'GELU: Smooth, Probabilistic gating',
fontsize=12, fontweight='bold')
ax5.legend()
ax5.grid(True, alpha=0.3)
ax5.set_xlim(-3, 3)
ax5.set_ylim(-0.5, 3)
# 6. 深度扩展性对比（训练稳定性模拟）
ax6 = axes[1, 2]
depths = [6, 12, 24, 48, 96]
# 模拟不同深度下的收敛损失（基于文献趋势）
pre_ln_loss = [2.5, 2.3, 2.1, 2.0, 1.95] # Pre-LN深度扩展性好
post_ln_loss = [2.4, 2.2, 2.5, 3.5, 4.5] # Post-LN深层不稳定
ax6.plot(depths, pre_ln_loss, 'o-', label='Pre-LN',
color='steelblue', linewidth=2, markersize=8)
ax6.plot(depths, post_ln_loss, 's-', label='Post-LN',
color='coral', linewidth=2, markersize=8)
ax6.axvline(x=24, color='red', linestyle='--', alpha=0.5)
ax6.text(26, 3.0, 'Instability\nThreshold', fontsize=9, color='red')
ax6.set_xlabel('Model Depth (Layers)', fontsize=11)
ax6.set_ylabel('Final Training Loss', fontsize=11)
ax6.set_title('Scalability: Depth vs Training Stability\n'
'Post-LN unstable beyond 24 layers',
fontsize=12, fontweight='bold')
ax6.legend()
ax6.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('transformer_architecture_analysis.png', dpi=150, bbox_inches='tight')
print("Visualization saved to transformer_architecture_analysis.png")
plt.show()
if __name__ == "__main__":
visualize_architecture_comparison()10.2.5 高效注意力变体：Sparse Transformer的稀疏模式、Linear Attention的核技巧、Flash Attention的IO感知分块与Softmax在线统计量计算
针对标准注意力机制的二次方计算复杂度与显存占用问题，高效注意力变体通过算法与系统协同优化实现扩展性提升。稀疏Transformer识别注意力矩阵的固有稀疏结构，采用固定模式如跨步稀疏保留固定间隔位置的键查询交互，或采用局部窗口约束注意力于邻域范围，以及固定列模式允许全局关注特定代表位置。这种结构化稀疏将计算复杂度降至线性或对数级，同时保持对长程依赖的捕获能力。
线性注意力通过核技巧重构计算顺序，利用矩阵乘法的结合律将复杂度从序列长度的平方降至线性。该方法将Softmax指数运算分解为特征映射的点积，先聚合键值乘积再与查询相乘，避免了显式计算完整的注意力矩阵。核函数选择影响表达能力与近似精度，常见实现采用逐元素指数或多项式展开模拟原始Softmax行为。
Flash Attention则从IO复杂度角度重新审视注意力计算，通过分块策略将输入序列划分为适配GPU共享内存容量的块，在片上完成局部注意力计算。算法通过在线统计量跟踪维护Softmax归一化因子，仅传递归一化后的输出与统计量而非完整的中间矩阵，将显存占用从平方降至线性。这种计算重排序显著减少高带宽内存访问，在现代加速器上实现数倍乃至数量级的墙钟时间加速。
脚本：efficient_attention_variants.py
Python
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脚本内容：高效注意力变体实现（Sparse、Linear、Flash Attention原理演示）
使用方式：直接运行 python efficient_attention_variants.py
依赖包：torch, numpy, matplotlib
功能说明：
1. 实现Sparse Transformer的多种稀疏模式（Strided、Fixed、Local）
2. 实现Linear Attention的核技巧（计算顺序重排）
3. 模拟Flash Attention的IO感知分块与在线Softmax统计量
4. 对比计算复杂度与内存占用的理论/实际差异
"""
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from typing import Optional, Tuple
import math
class SparseAttentionPattern:
"""Sparse Transformer的稀疏模式实现"""
def __init__(self, seq_len: int, pattern: str = 'strided',
stride: int = 4, local_window: int = 8):
self.seq_len = seq_len
self.pattern = pattern
self.stride = stride
self.local_window = local_window
def create_mask(self) -> torch.Tensor:
"""创建稀疏注意力掩码（1表示允许关注，0表示屏蔽）"""
mask = torch.zeros(self.seq_len, self.seq_len)
if self.pattern == 'strided':
# 跨步稀疏：只关注每隔stride个位置
for i in range(self.seq_len):
# 关注所有stride的倍数位置
strided_indices = torch.arange(0, self.seq_len, self.stride)
mask[i, strided_indices] = 1
# 总是关注自己
mask[i, i] = 1
elif self.pattern == 'local':
# 局部窗口：只关注前后local_window个位置
for i in range(self.seq_len):
start = max(0, i - self.local_window)
end = min(self.seq_len, i + self.local_window + 1)
mask[i, start:end] = 1
elif self.pattern == 'fixed':
# 固定列：选择少数代表列（如0, stride, 2*stride...）
fixed_cols = list(range(0, self.seq_len, self.stride))
mask[:, fixed_cols] = 1
# 添加对角线（局部）
mask += torch.eye(self.seq_len)
mask = (mask > 0).float()
elif self.pattern == 'block':
# 分块稀疏：将序列分成块，块内全连接，块间稀疏
block_size = self.stride
num_blocks = self.seq_len // block_size
for b in range(num_blocks):
start = b * block_size
end = min(start + block_size, self.seq_len)
mask[start:end, start:end] = 1
# 块间连接：每个块关注第一个块
mask[start:end, :block_size] = 1
return mask
def get_sparse_indices(self) -> Tuple[torch.Tensor, torch.Tensor]:
"""获取稀疏索引用于高效计算"""
mask = self.create_mask()
row_indices, col_indices = torch.where(mask > 0)
return row_indices, col_indices
class LinearAttention(nn.Module):
"""线性注意力：通过核技巧实现O(n)复杂度"""
def __init__(self, dim: int, num_heads: int = 8, eps: float = 1e-6):
super().__init__()
self.dim = dim
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.eps = eps
# 投影层
self.W_q = nn.Linear(dim, dim, bias=False)
self.W_k = nn.Linear(dim, dim, bias=False)
self.W_v = nn.Linear(dim, dim, bias=False)
self.W_o = nn.Linear(dim, dim, bias=False)
# 核特征映射（模拟Softmax的指数函数）
self.feature_dim = self.head_dim # 可以更大以提升表达能力
def _kernel_feature_map(self, x: torch.Tensor) -> torch.Tensor:
"""
核特征映射：将向量映射到高维空间以近似点积
使用逐元素指数（Eluder核）或正交随机特征
"""
# 方案1：逐元素指数（exp）模拟Softmax
# 添加非线性变换后指数化
x = torch.nn.functional.elu(x) + 1 # ELU+1保证正数
return x
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
线性注意力前向传播
复杂度从O(n^2 d)降至O(n d^2)
"""
batch_size, seq_len, _ = x.size()
# 投影并分头
Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
# Q,K,V: [batch, num_heads, seq_len, head_dim]
# 核特征映射
Q_prime = self._kernel_feature_map(Q) # [batch, heads, seq, features]
K_prime = self._kernel_feature_map(K)
# 关键：改变计算顺序（核技巧）
# 标准：Q @ K.T @ V -> O(n^2)
# 线性：(Q @ (K.T @ V)) -> O(n)
# 步骤1：聚合K和V（KV聚合）
# [batch, heads, features, features]
KV = torch.matmul(K_prime.transpose(-2, -1), V)
# 步骤2：与Q相乘
# [batch, heads, seq, features]
Z = torch.matmul(Q_prime, KV)
# 归一化（模拟Softmax的分母）
# [batch, heads, seq]
denominator = torch.matmul(Q_prime, K_prime.sum(dim=-2).unsqueeze(-1)) + self.eps
output = Z / denominator
# 重塑并输出投影
output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.dim)
output = self.W_o(output)
return output
def compute_complexity(self, seq_len: int) -> Tuple[int, int]:
"""返回时间复杂度和空间复杂度（相对单位）"""
# 标准注意力：O(n^2 d)，空间O(n^2)
standard_time = seq_len * seq_len * self.head_dim * self.num_heads
standard_space = seq_len * seq_len
# 线性注意力：O(n d^2)，空间O(d^2)（忽略n）
linear_time = seq_len * self.head_dim * self.head_dim * self.num_heads
linear_space = self.head_dim * self.head_dim * self.num_heads
return standard_time, linear_time, standard_space, linear_space
class FlashAttentionSimulator:
"""Flash Attention的原理模拟（分块与在线Softmax）"""
def __init__(self, dim: int, block_size: int = 64, num_heads: int = 8):
self.dim = dim
self.block_size = block_size
self.num_heads = num_heads
self.head_dim = dim // num_heads
def _online_softmax(self, local_scores: torch.Tensor,
prev_max: torch.Tensor,
prev_sum: torch.Tensor):
"""
在线Softmax计算（Flash Attention核心）
增量更新max和sum，避免存储完整attention矩阵
"""
# 当前块的局部max
local_max = torch.max(local_scores, dim=-1, keepdim=True)[0]
# 更新全局max
new_max = torch.max(prev_max, local_max)
# 重缩放之前的sum
prev_sum_scaled = prev_sum * torch.exp(prev_max - new_max)
# 当前块的sum（按新的max重缩放）
local_sum = torch.exp(local_scores - new_max).sum(dim=-1, keepdim=True)
new_sum = prev_sum_scaled + local_sum
return new_max, new_sum
def simulate_forward(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor):
"""
模拟Flash Attention的分块计算流程
展示IO感知和在线统计量机制
"""
batch_size, num_heads, seq_len, head_dim = Q.size()
# 初始化输出和在线统计量
output = torch.zeros_like(Q)
running_max = torch.full((batch_size, num_heads, seq_len, 1), float('-inf'))
running_sum = torch.zeros(batch_size, num_heads, seq_len, 1)
num_blocks = math.ceil(seq_len / self.block_size)
# 模拟分块计算
for block_idx in range(num_blocks):
start = block_idx * self.block_size
end = min(start + self.block_size, seq_len)
# 加载K,V块到SRAM（快速内存）
K_block = K[:, :, start:end, :]
V_block = V[:, :, start:end, :]
# 计算当前块的注意力分数（在SRAM内完成）
scores_block = torch.matmul(Q, K_block.transpose(-2, -1)) / math.sqrt(head_dim)
# 在线Softmax更新
running_max, running_sum = self._online_softmax(
scores_block, running_max, running_sum
)
# 计算当前块对输出的贡献（重缩放后）
# 实际上Flash Attention会累积输出而非直接存储attention
# 这里模拟核心思想
exp_scores = torch.exp(scores_block - running_max)
weighted_v = torch.matmul(exp_scores, V_block)
# 更新输出（后续块会继续累加并重缩放）
# 实际实现更复杂，这里展示概念
output += weighted_v / running_sum
return output
def analyze_io_cost(self, seq_len: int, d_model: int):
"""分析标准注意力和Flash Attention的HBM访问次数"""
# HBM（高带宽内存）访问是瓶颈
# 标准注意力
# 1. 读取Q,K,V：3 * seq * d
# 2. 写入S矩阵：seq * seq
# 3. 读取S矩阵计算softmax：seq * seq
# 4. 写入P矩阵（attention）：seq * seq
# 5. 读取P和V计算输出：seq * seq + seq * d
# 6. 写入输出：seq * d
standard_hbm = 3*seq_len*d_model + 2*seq_len*seq_len + seq_len*seq_len + seq_len*seq_len + seq_len*d_model + seq_len*d_model
standard_hbm_accesses = 3 + 2*seq_len*seq_len/d_model + seq_len # 以读取Q为单位简化
# Flash Attention
# 分块后，S和P矩阵不写入HBM，仅在SRAM计算
# 只需要读取Q,K,V一次，写入输出一次
# 但分块增加了多次读取K,V的开销
block_size = self.block_size
num_blocks = seq_len / block_size
flash_hbm_accesses = 3 + 2 * num_blocks # 简化估计
return standard_hbm_accesses, flash_hbm_accesses
def visualize_efficient_attention():
"""可视化高效注意力变体的特性"""
seq_len = 512
d_model = 256
num_heads = 8
fig, axes = plt.subplots(2, 3, figsize=(16, 10))
# 1. 稀疏注意力模式对比
ax1 = axes[0, 0]
patterns = ['strided', 'local', 'fixed', 'block']
titles = ['Strided', 'Local Window', 'Fixed Columns', 'Block Sparse']
# 创建子图展示不同模式
fig_sparse, axes_sparse = plt.subplots(2, 2, figsize=(10, 10))
for ax_sparse, pattern, title in zip(axes_sparse.flat, patterns, titles):
sparse_attn = SparseAttentionPattern(seq_len=64, pattern=pattern, stride=8, local_window=8)
mask = sparse_attn.create_mask().numpy()
im = ax_sparse.imshow(mask, cmap='Blues', interpolation='nearest', aspect='equal')
ax_sparse.set_title(title, fontsize=11, fontweight='bold')
ax_sparse.set_xlabel('Key Position', fontsize=9)
ax_sparse.set_ylabel('Query Position', fontsize=9)
# 计算稀疏度
sparsity = 1 - mask.sum() / mask.size
ax_sparse.text(2, 5, f'Sparsity: {sparsity:.1%}',
color='red', fontsize=9, fontweight='bold',
bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
plt.colorbar(im, ax=ax_sparse, fraction=0.046, pad=0.04)
plt.tight_layout()
plt.savefig('sparse_patterns.png', dpi=150, bbox_inches='tight')
print("Sparse patterns saved to sparse_patterns.png")
# 在主图中显示统计对比
sparsities = []
for pattern in patterns:
sparse_obj = SparseAttentionPattern(seq_len=seq_len, pattern=pattern)
mask = sparse_obj.create_mask()
sparsity = 1 - mask.sum().item() / (seq_len * seq_len)
sparsities.append(sparsity)
# 绘制稀疏度对比
colors = ['steelblue', 'coral', 'forestgreen', 'purple']
bars = ax1.bar(patterns, [s*100 for s in sparsities], color=colors, alpha=0.8)
ax1.set_ylabel('Sparsity %', fontsize=11)
ax1.set_title('Sparse Attention Patterns\nReduction in Computation',
fontsize=12, fontweight='bold')
ax1.set_ylim(0, 100)
# 添加数值标签
for bar, sparsity in zip(bars, sparsities):
height = bar.get_height()
ax1.text(bar.get_x() + bar.get_width()/2., height + 1,
f'{sparsity:.1%}', ha='center', va='bottom', fontsize=10)
# 2. 线性注意力的核技巧示意图
ax2 = axes[0, 1]
# 计算复杂度对比
lengths = [64, 128, 256, 512, 1024, 2048, 4096]
standard_complexity = [n**2 for n in lengths]
linear_complexity = [n for n in lengths] # 实际上是O(n d^2)，这里简化为相对n
ax2.plot(lengths, standard_complexity, 'o-', label='Standard Attention O(n²)',
color='crimson', linewidth=2, markersize=8)
ax2.plot(lengths, linear_complexity, 's-', label='Linear Attention O(n)',
color='forestgreen', linewidth=2, markersize=8)
# 标记关键长度
ax2.axvline(x=1024, color='orange', linestyle='--', alpha=0.5)
ax2.text(1100, 100000, 'Typical Long\nContext Threshold', fontsize=9)
ax2.set_xlabel('Sequence Length', fontsize=11)
ax2.set_ylabel('Relative Time Complexity', fontsize=11)
ax2.set_title('Computational Complexity Scaling\nKernel Trick Reorders Computation',
fontsize=12, fontweight='bold')
ax2.set_yscale('log')
ax2.set_xscale('log')
ax2.legend()
ax2.grid(True, alpha=0.3)
# 3. Flash Attention分块与内存层次
ax3 = axes[0, 2]
# 模拟GPU内存层次
memory_types = ['HBM\n(High BW)', 'SRAM\n(Shared Mem)', 'Registers']
capacities = [40, 0.1, 0.001] # MB级别（相对）
speeds = [1.5, 19, 100] # TB/s级别（相对）
x_pos = np.arange(len(memory_types))
width = 0.35
bars1 = ax3.bar(x_pos - width/2, capacities, width, label='Capacity (MB)',
color='steelblue', alpha=0.8)
ax3_twin = ax3.twinx()
bars2 = ax3_twin.bar(x_pos + width/2, speeds, width, label='Bandwidth (TB/s)',
color='coral', alpha=0.8)
ax3.set_ylabel('Capacity (MB, log)', fontsize=11)
ax3_twin.set_ylabel('Bandwidth (TB/s)', fontsize=11)
ax3.set_title('GPU Memory Hierarchy\nFlash Attention Optimizes Data Movement',
fontsize=12, fontweight='bold')
ax3.set_xticks(x_pos)
ax3.set_xticklabels(memory_types)
ax3.set_yscale('log')
# 添加图例
lines1, labels1 = ax3.get_legend_handles_labels()
lines2, labels2 = ax3_twin.get_legend_handles_labels()
ax3.legend(lines1 + lines2, labels1 + labels2, loc='upper right')
# 添加Flash Attention优化说明
ax3.annotate('Flash Attention keeps\ncomputations in SRAM',
xy=(1, 0.1), xytext=(1.5, 10),
arrowprops=dict(arrowstyle='->', color='red', lw=2),
fontsize=10, color='red', fontweight='bold',
bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.7))
# 4. 在线Softmax统计量示意
ax4 = axes[1, 0]
# 模拟分块Softmax的计算过程
blocks = ['Block 1', 'Block 2', 'Block 3', 'Block 4']
local_max = [2.5, 3.0, 2.8, 3.2]
running_max = [2.5, 3.0, 3.0, 3.2]
memory_usage = [1.0, 1.0, 1.0, 1.0] # Flash: O(1) 每块
x = np.arange(len(blocks))
width = 0.25
ax4.bar(x - width, local_max, width, label='Local Max', color='lightblue', alpha=0.8)
ax4.bar(x, running_max, width, label='Running Global Max', color='darkblue', alpha=0.8)
ax4.bar(x + width, memory_usage, width, label='Memory (relative)', color='green', alpha=0.6)
ax4.set_ylabel('Score Value / Memory', fontsize=11)
ax4.set_title('Online Softmax: Block-wise Statistics\nNo Materialized Attention Matrix',
fontsize=12, fontweight='bold')
ax4.set_xticks(x)
ax4.set_xticklabels(blocks)
ax4.legend()
ax4.grid(True, alpha=0.3, axis='y')
# 添加流程箭头
for i in range(len(blocks)-1):
ax4.annotate('', xy=(i+0.6, running_max[i]), xytext=(i+0.1, running_max[i]),
arrowprops=dict(arrowstyle='->', color='red', lw=2))
ax4.text(i+0.35, running_max[i]+0.1, 'update', fontsize=8, color='red')
# 5. HBM访问次数对比（IO分析）
ax5 = axes[1, 1]
lengths_io = [512, 1024, 2048, 4096, 8192]
flash_sim = FlashAttentionSimulator(d_model, block_size=64)
standard_io = []
flash_io = []
for length in lengths_io:
std, flash = flash_sim.analyze_io_cost(length, d_model)
standard_io.append(std)
flash_io.append(flash)
ax5.plot(lengths_io, standard_io, 'o--', label='Standard Attention',
color='crimson', linewidth=2, markersize=8)
ax5.plot(lengths_io, flash_io, 's-', label='Flash Attention',
color='forestgreen', linewidth=2, markersize=8)
# 计算加速比
speedup = [s/f for s, f in zip(standard_io, flash_io)]
for i, (length, sp) in enumerate(zip(lengths_io, speedup)):
ax5.annotate(f'{sp:.1f}×', xy=(length, flash_io[i]),
xytext=(length, flash_io[i]-2),
fontsize=9, ha='center', color='darkgreen', fontweight='bold')
ax5.set_xlabel('Sequence Length', fontsize=11)
ax5.set_ylabel('HBM Accesses (relative)', fontsize=11)
ax5.set_title('IO-Aware Optimization: HBM Access Reduction\n'
'Wall-clock speedup correlates with IO reduction',
fontsize=12, fontweight='bold')
ax5.legend()
ax5.grid(True, alpha=0.3)
# 6. 综合效率对比（实际wall-clock时间模拟）
ax6 = axes[1, 2]
methods = ['Standard', 'Sparse\n(Strided)', 'Linear\n(Kernel)', 'Flash\n(IO-aware)']
# 模拟不同序列长度下的相对时间（归一化到标准Attention在1024长度的时间）
seq_lens_comp = [512, 1024, 2048, 4096]
# 相对时间数据（基于文献典型值）
standard_times = [0.25, 1.0, 4.0, 16.0] # O(n^2)
sparse_times = [0.15, 0.4, 1.2, 3.6] # O(n sqrt(n))或O(n log n)
linear_times = [0.3, 0.6, 1.2, 2.4] # O(n)但常数较大
flash_times = [0.2, 0.35, 0.7, 1.4] # O(n^2)但IO优化
x = np.arange(len(seq_lens_comp))
width = 0.2
ax6.bar(x - 1.5*width, standard_times, width, label='Standard', color='crimson', alpha=0.8)
ax6.bar(x - 0.5*width, sparse_times, width, label='Sparse', color='steelblue', alpha=0.8)
ax6.bar(x + 0.5*width, linear_times, width, label='Linear', color='forestgreen', alpha=0.8)
ax6.bar(x + 1.5*width, flash_times, width, label='Flash', color='purple', alpha=0.8)
ax6.set_xlabel('Sequence Length', fontsize=11)
ax6.set_ylabel('Relative Wall-Clock Time', fontsize=11)
ax6.set_title('End-to-End Efficiency Comparison\n'
'Flash Attention wins in practice despite O(n^2)',
fontsize=12, fontweight='bold')
ax6.set_xticks(x)
ax6.set_xticklabels(seq_lens_comp)
ax6.legend()
ax6.grid(True, alpha=0.3, axis='y')
ax6.set_yscale('log')
plt.tight_layout()
plt.savefig('efficient_attention_comparison.png', dpi=150, bbox_inches='tight')
print("Main visualization saved to efficient_attention_comparison.png")
plt.show()
if __name__ == "__main__":
visualize_efficient_attention()以上交付物包含8个独立脚本，涵盖第10章全部技术内容。每个脚本均包含深度实现与专门可视化，可直接复制使用。技术内容参考了Vaswani et al. (2017), Bahdanau et al. (2015), Luong et al. (2015), Child et al. (Sparse Transformer), Katharopoulos et al. (Linear Attention), Dao et al. (Flash Attention), Su et al. (RoPE), Xiong et al. (Pre-LN/Post-LN分析) 等经典与前沿工作。
3355
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
