线性注意力架构演进：从 Transformer 瓶颈到高效推理

import torch
import math

def standard_attention(Q, K, V, mask=None):
    """标准 Scaled Dot-Product Attention"""
    d = Q.shape[-1]
    # 第一步: QK^T 矩阵乘法，复杂度 O(n²d)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d)
    # 第二步: Softmax 归一化，复杂度 O(n²)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    weights = torch.softmax(scores, dim=-1)
    # 第三步: 权重 × V，复杂度 O(n²d)
    output = torch.matmul(weights, V)
    return output, weights

def analyze_kv_cache_growth(seq_length, d_model, num_heads, num_layers):
    """分析 KV Cache 随序列长度的增长"""
    d_head = d_model // num_heads
    # 每层的 KV Cache 大小: 2 × seq_length × num_heads × d_head
    per_layer_cache = 2 * seq_length * num_heads * d_head * 4  # float32
    total_cache = per_layer_cache * num_layers
    print(f"序列长度: {seq_length:,}")
    print(f"单层 KV Cache: {per_layer_cache / 1024**2:.1f} MB")
    print(f"总 KV Cache ({num_layers} 层): {total_cache / 1024**2:.1f} MB")
    return total_cache

# 对比不同序列长度下的 KV Cache 大小
for seq_len in [4096, 32768, 131072]:
    print("---")
    analyze_kv_cache_growth(seq_len, d_model=4096, num_heads=32, num_layers=32)

import torch
import torch.nn.functional as F

class LinearAttention(torch.nn.Module):
    """线性注意力实现：核方法 + 递归推理"""
    
    def __init__(self, d_model, d_head):
        super().__init__()
        self.d_model = d_model
        self.d_head = d_head
        self.num_heads = d_model // d_head
        self.q_proj = torch.nn.Linear(d_model, d_model)
        self.k_proj = torch.nn.Linear(d_model, d_model)
        self.v_proj = torch.nn.Linear(d_model, d_model)
        self.out_proj = torch.nn.Linear(d_model, d_model)
    
    def feature_map(self, x):
        """特征映射：使用 elu + 1 作为核函数"""
        x = F.elu(x) + 1.0
        return x
    
    def forward_parallel(self, Q, K, V):
        """并行训练模式：利用矩阵乘法结合律"""
        # 形状: (batch, heads, seq, d_head)
        Q = self.feature_map(Q)
        K = self.feature_map(K)
        # 关键技巧: 先算 K^T @ V (d_head, d_head)，再左乘 Q
        # 避免生成 (seq, seq) 的注意力矩阵
        KV = torch.einsum("nhld,nhlm->nhdm", K, V)
        Z = torch.einsum("nhld,nhld->nhl", Q, K.sum(dim=-2) + 1e-6)
        output = torch.einsum("nhld,nhdm->nhlm", Q, KV) / Z.unsqueeze(-1)
        return output
    
    def forward_recurrent(self, x_seq):
        """递归推理模式：固定大小状态矩阵"""
        batch_size, seq_len, _ = x_seq.shape
        Q = self.q_proj(x_seq)
        K = self.k_proj(x_seq)
        V = self.v_proj(x_seq)
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_head)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_head)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_head)
        
        outputs = []
        # 状态矩阵: (batch, heads, d_head, d_head)，大小固定，不随序列增长
        state = torch.zeros(batch_size, self.num_heads, self.d_head, self.d_head, device=x_seq.device)
        for t in range(seq_len):
            kt = K[:, t, :, :].unsqueeze(-1)  # (batch, heads, d_head, 1)
            vt = V[:, t, :, :].unsqueeze(-2)  # (batch, heads, 1, d_head)
            state = state + torch.matmul(kt, vt)  # 递归更新
            qt = Q[:, t, :, :]  # (batch, heads, d_head)
            out_t = torch.matmul(qt.unsqueeze(-2), state).squeeze(-2)  # (batch, heads, d_head)
            outputs.append(out_t)
        
        output = torch.stack(outputs, dim=1)
        return output

# 四种线性注意力架构的核心计算对比
import torch

# 1. RWKV 时间混合（简化版）
def rwkv_time_mix(x, w, k, v, r, state):
    """RWKV: 接受度加权键值混合"""
    # w: 时间衰减权重, k: 键, v: 值, r: 接受度
    # state: 递归状态 (num_heads, d_head)
    state = state * w + k * v  # 递归更新
    output = r * state  # 接受度加权
    return output, state

# 2. Mamba 选择性扫描（简化版）
def mamba_selective_scan(x, delta, A, B, C):
    """Mamba: 选择性状态空间扫描"""
    # delta: 输入依赖的步长, A: 状态矩阵
    # B, C: 输入输出投影
    batch, seq, d_state = x.shape
    state = torch.zeros(batch, d_state)
    outputs = []
    for t in range(seq):
        # 选择性更新：根据输入动态调整
        dt = delta[:, t, :]  # 输入依赖的步长
        state = state * torch.exp(dt * A) + dt * B[:, t, :] * x[:, t, :]
        out = torch.matmul(state, C[:, t, :].transpose(-1, -2))
        outputs.append(out)
    return torch.stack(outputs, dim=1)

# 3. RetNet 多尺度保留（简化版）
def retnet_multi_scale(q, k, v, decay_factors):
    """RetNet: 多尺度保留"""
    # decay_factors: 不同层的衰减因子列表
    # 并行模式: 使用保留矩阵
    batch, heads, seq, d = q.shape
    # 构建保留矩阵 (seq, seq)，但利用结构化特性避免 O(n²)
    decay_matrix = torch.zeros(seq, seq, device=q.device)
    for i in range(seq):
        for j in range(i + 1):
            decay_matrix[i, j] = decay_factors[0] ** (i - j)
    scores = torch.matmul(q, k.transpose(-2, -1)) * decay_matrix.unsqueeze(0).unsqueeze(0)
    return torch.matmul(scores, v)

# 4. Gated DeltaNet（简化版）
def gated_deltanet(q, k, v, gate, state, alpha=0.1):
    """Gated DeltaNet: 门控 Delta 规则"""
    # gate: 门控信号, alpha: 学习率
    # Delta 规则: 只更新差异大的部分
    delta = torch.matmul(k.unsqueeze(-1), v.unsqueeze(-2))  # 外积
    # 门控控制更新幅度
    gate_signal = gate.sigmoid()
    state = state + alpha * gate_signal.unsqueeze(-1) * delta
    output = torch.matmul(q.unsqueeze(-2), state).squeeze(-2)
    return output, state

import torch
import torch.nn as nn
import torch.nn.functional as F

class GatedDeltaNetBlock(nn.Module):
    """Gated DeltaNet 块：门控 + Delta 规则的完整实现"""
    
    def __init__(self, d_model, d_state=64, alpha=0.1):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.alpha = alpha
        
        # 门控投影
        self.gate_proj = nn.Linear(d_model, d_model)
        # 键值投影（带特征映射）
        self.q_proj = nn.Linear(d_model, d_state)
        self.k_proj = nn.Linear(d_model, d_state)
        self.v_proj = nn.Linear(d_model, d_state)
        # 输出投影
        self.out_proj = nn.Linear(d_state, d_model)
        # 预归一化
        self.norm = nn.LayerNorm(d_model)
    
    def feature_map(self, x):
        """特征映射：elu + 1"""
        return F.elu(x) + 1.0
    
    def forward_training(self, x):
        """训练模式：并行计算"""
        x = self.norm(x)
        gate = torch.sigmoid(self.gate_proj(x))  # (batch, seq, d_model)
        Q = self.feature_map(self.q_proj(x))
        K = self.feature_map(self.k_proj(x))
        V = self.v_proj(x)
        
        # Delta 规则的并行实现
        # 状态矩阵序列: (batch, seq, d_state, d_state)
        batch, seq, d_state = Q.shape
        # 累积状态: KV 矩阵
        KV_cumsum = torch.zeros(batch, seq, d_state, d_state, device=x.device)
        for t in range(seq):
            delta = torch.einsum("bd,bd->bdd", K[:, t, :], V[:, t, :])
            if t > 0:
                KV_cumsum[:, t, :, :] = KV_cumsum[:, t-1, :, :] + self.alpha * gate[:, t, :].unsqueeze(-1) * delta
            else:
                KV_cumsum[:, t, :, :] = self.alpha * gate[:, t, :].unsqueeze(-1) * delta
        
        # 输出: Q 点乘累积状态
        output = torch.einsum("bd,bdd->bd", Q, KV_cumsum)
        return self.out_proj(output)
    
    def step(self, x_t, state):
        """推理步骤：递归更新"""
        x_t = self.norm(x_t)
        gate = torch.sigmoid(self.gate_proj(x_t))
        q_t = self.feature_map(self.q_proj(x_t))
        k_t = self.feature_map(self.k_proj(x_t))
        v_t = self.v_proj(x_t)
        
        # Delta 更新
        delta = torch.einsum("bd,bd->bdd", k_t, v_t)
        state = state + self.alpha * gate.unsqueeze(-1) * delta
        
        # 输出
        output = torch.einsum("bd,bdd->bd", q_t, state)
        output = self.out_proj(output)
        return output, state

# 线性注意力训练的稳定配置
from transformers import get_cosine_schedule_with_warmup

def create_stable_training_config(model, train_dataset, args):
    """为线性注意力模型创建稳定训练配置"""
    
    # 训练参数
    training_args = {
        "learning_rate": 1e-4,           # 比 Transformer 小 3 倍
        "warmup_ratio": 0.10,            # 更长的 warmup
        "per_device_train_batch_size": 32,
        "gradient_accumulation_steps": 8, # 有效 batch size = 256
        "gradient_clip_val": 1.0,         # 严格的梯度裁剪
        "max_grad_norm": 1.0,
        "weight_decay": 0.1,
        "adam_beta1": 0.9,
        "adam_beta2": 0.95,              # 更平滑的二阶矩估计
        "adam_epsilon": 1e-8,
        "fp16": False,                    # 线性注意力建议 bf16
        "bf16": True,
        "dataloader_num_workers": 4,
    }
    
    # 状态矩阵归一化回调
    class StateNormCallback:
        def on_step_end(self, args, state, model, **kwargs):
            for module in model.modules():
                if hasattr(module, 'state_matrix'):
                    # 对状态矩阵做归一化防止数值溢出
                    with torch.no_grad():
                        norm = module.state_matrix.norm(p=2, dim=(-1, -2), keepdim=True)
                        norm = torch.clamp(norm, min=1e-6)
                        module.state_matrix = module.state_matrix / norm * 10.0
    
    return training_args, StateNormCallback()

# 方案一: 使用 RWKV 进行推理
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载预训练 RWKV 模型
model_name = "rwkv-4-world-1.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)

# RWKV 推理: 不需要 KV Cache 管理
prompt = "线性注意力的核心思想是"
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(
    inputs["input_ids"],
    max_new_tokens=200,
    do_sample=False,  # 贪婪解码
    use_cache=True,   # RWKV 内部状态缓存
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# 方案二: 使用 Mamba 模块替换 Transformer 层
import torch
from mamba_ssm import Mamba

class MambaBlock(torch.nn.Module):
    """用 Mamba 块替换 Transformer 的注意力层"""
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        self.norm = torch.nn.LayerNorm(d_model)
        self.mamba = Mamba(
            d_model=d_model,
            d_state=d_state,
            d_conv=d_conv,
            expand=expand,
        )
    
    def forward(self, x):
        # x: (batch, seq_len, d_model)
        x = x + self.mamba(self.norm(x))
        return x

# 使用示例: 替换 Transformer 编码器的注意力层
class MambaEncoder(torch.nn.Module):
    def __init__(self, d_model, n_layers, d_state=16):
        super().__init__()
        self.layers = torch.nn.ModuleList([
            MambaBlock(d_model, d_state=d_state)
            for _ in range(n_layers)
        ])
        self.norm = torch.nn.LayerNorm(d_model)
    
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return self.norm(x)

# 测试
x = torch.randn(2, 1000, 512)  # (batch, seq, d_model)
encoder = MambaEncoder(d_model=512, n_layers=6, d_state=16)
output = encoder(x)
print(f"输出形状: {output.shape}")  # (2, 1000, 512)

# 方案三: 从零实现线性注意力并替换到现有模型
import torch
import torch.nn as nn
import torch.nn.functional as F

class LinearAttentionModule(nn.Module):
    """可插入的线性注意力模块"""
    
    def __init__(self, d_model, num_heads, feature_map="elu"):
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.d_head = d_model // num_heads
        self.feature_map_name = feature_map
        
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
    
    def get_feature_map(self, x):
        if self.feature_map_name == "elu":
            return F.elu(x) + 1.0
        elif self.feature_map_name == "relu":
            return F.relu(x) + 1e-5
        elif self.feature_map_name == "sigmoid":
            return F.sigmoid(x) + 1e-5
        else:
            raise ValueError(f"未知的特征映射: {self.feature_map_name}")
    
    def forward(self, x, mask=None):
        batch, seq, _ = x.shape
        Q = self.q_proj(x).view(batch, seq, self.num_heads, self.d_head).transpose(1, 2)
        K = self.k_proj(x).view(batch, seq, self.num_heads, self.d_head).transpose(1, 2)
        V = self.v_proj(x).view(batch, seq, self.num_heads, self.d_head).transpose(1, 2)
        
        Q = self.get_feature_map(Q)
        K = self.get_feature_map(K)
        
        # 线性注意力核心: 利用结合律
        # O(n*d²) 而不是 O(n²*d)
        KV = torch.einsum("bhld,bhlm->bhdm", K, V)
        Z = 1.0 / (torch.einsum("bhld,bhd->bhl", Q, K.sum(dim=2)) + 1e-6)
        out = torch.einsum("bhld,bhdm->bhl", Q, KV) * Z.unsqueeze(-1)
        
        out = out.transpose(1, 2).contiguous().view(batch, seq, -1)
        return self.out_proj(out)