大模型应用：稀疏注意力 vs 滑动窗口：大模型扩窗技术完全解析.58

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
# ---------------------- 核心函数 ----------------------
def sliding_window_attention(q, k, v, window_size=4):
    """滑动窗口注意力（带关注矩阵记录）"""
    seq_len, d = q.shape
    output = torch.zeros_like(v)
    # 记录关注矩阵（seq_len×seq_len，1=关注，0=不关注）
    attention_mask = torch.zeros((seq_len, seq_len))
    
    for i in range(seq_len):
        # 计算窗口范围：当前位置i向前看window_size个位置
        start = max(0, i - window_size + 1)
        end = i + 1
        # 标记关注的位置（从start到i）
        attention_mask[i, start:end] = 1
        
        q_i = q[i:i+1, :]
        k_window = k[start:end, :]
        v_window = v[start:end, :]
        
        scores = torch.matmul(q_i, k_window.T) / torch.sqrt(torch.tensor(d))
        weights = F.softmax(scores, dim=-1)
        output[i, :] = torch.matmul(weights, v_window)
    
    return output, attention_mask

def sparse_attention(q, k, v, window_size=2, stride=3):
    """稀疏注意力（带关注矩阵记录）"""
    seq_len, d = q.shape
    output = torch.zeros_like(v)
    attention_mask = torch.zeros((seq_len, seq_len))
    
    for i in range(seq_len):
        # 1. 局部窗口：当前位置i向前看window_size个位置
        start = max(0, i - window_size + 1)
        end = i + 1
        local_j = list(range(start, end))
        
        # 2. 跨步位置：从当前位置i向前按stride步长跳跃
        stride_j = []
        j = i - stride
        while j >= 0:
            stride_j.append(j)
            j -= stride
        
        # 3. 全局位置（首尾）
        global_j = [0, seq_len-1] if seq_len > 1 else [0]
        
        # 合并去重
        all_j = list(set(local_j + stride_j + global_j))
        all_j.sort()  # 排序保证顺序
        # 标记关注的位置
        attention_mask[i, all_j] = 1
        
        q_i = q[i:i+1, :]
        k_sparse = k[all_j, :]
        v_sparse = v[all_j, :]
        
        scores = torch.matmul(q_i, k_sparse.T) / torch.sqrt(torch.tensor(d))
        weights = F.softmax(scores, dim=-1)
        output[i, :] = torch.matmul(weights, v_sparse)
    
    return output, attention_mask

# ---------------------- 可视化函数 ----------------------
def plot_attention_comparison(sw_mask, sparse_mask, seq_len):
    """绘制注意力关注矩阵对比图（生成图片）"""
    # 创建画布（1行3列，对比原始/滑动窗口/稀疏注意力）
    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(15, 5))
    
    # 1. 原始注意力（全关注）
    original_mask = torch.ones((seq_len, seq_len))
    im1 = ax1.imshow(original_mask.numpy(), cmap='Blues', vmin=0, vmax=1)
    ax1.set_title('原始注意力（全关注）', fontsize=12)
    ax1.set_xlabel('关注的位置j')
    ax1.set_ylabel('当前位置i')
    
    # 2. 滑动窗口注意力
    im2 = ax2.imshow(sw_mask.numpy(), cmap='Blues', vmin=0, vmax=1)
    ax2.set_title('滑动窗口注意力（W=3）', fontsize=12)
    ax2.set_xlabel('关注的位置j')
    
    # 3. 稀疏注意力
    im3 = ax3.imshow(sparse_mask.numpy(), cmap='Blues', vmin=0, vmax=1)
    ax3.set_title('稀疏注意力（W=2+S=3）', fontsize=12)
    ax3.set_xlabel('关注的位置j')
    
    # 统一刻度（方便对比）
    for ax in [ax1, ax2, ax3]:
        ax.set_xticks(range(seq_len))
        ax.set_yticks(range(seq_len))
        ax.set_xticklabels([f'{i+1}' for i in range(seq_len)])
        ax.set_yticklabels([f'{i+1}' for i in range(seq_len)])
        ax.set_xlabel('关注的位置j')
        ax.set_ylabel('当前位置i')
    
    # # 添加颜色条（放在右侧外部）
    # cbar = fig.colorbar(im1, ax=[ax1, ax2, ax3], shrink=0.8, pad=0.95, aspect=30)
    # cbar.set_label('关注状态（1=关注，0=不关注）', fontsize=10, rotation=270, labelpad=20)
    
    # 保存图片（新手可改路径）
    plt.tight_layout()
    plt.savefig('attention_comparison.png', dpi=300, bbox_inches='tight')
    print("图片已保存为：attention_comparison.png")
    plt.show()

# ---------------------- 测试+生成图片 ----------------------
if __name__ == "__main__":
    # 模拟输入：序列长度=8，向量维度=2
    seq_len = 8
    d = 2
    q = torch.randn(seq_len, d)
    k = torch.randn(seq_len, d)
    v = torch.randn(seq_len, d)
    
    # 计算滑动窗口注意力（窗口=3）
    sw_output, sw_mask = sliding_window_attention(q, k, v, window_size=3)
    # 计算稀疏注意力（窗口=2，跨步=3）
    sparse_output, sparse_mask = sparse_attention(q, k, v, window_size=2, stride=3)
    
    # 生成对比图
    plot_attention_comparison(sw_mask, sparse_mask, seq_len)
    
    # 打印关键信息（新手验证）
    print("滑动窗口注意力关注矩阵：")
    print(sw_mask.numpy().astype(int))  # 1=关注，0=不关注
    print("\n稀疏注意力关注矩阵：")
    print(sparse_mask.numpy().astype(int))

import matplotlib.pyplot as plt

# 配置字体（支持中文和符号）
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 序列设置
positions = list(range(1, 11))  # [1, 2, ..., 10]
n = len(positions)

# 1. 原始注意力：全部关注
original = [True] * n

# 2. 滑动窗口注意力 (W=3, i=5) → 关注 4,5,6
window = [False] * n
for idx in [4, 5, 6]:  # 位置 4,5,6（注意索引从1开始）
    window[idx - 1] = True

# 3. 稀疏注意力：关注 1,4,5,6,8,10（按你给出的结果）
sparse_set = {1, 4, 5, 6, 8, 10}
sparse = [pos in sparse_set for pos in positions]

# 所有行数据
rows = [
    {
        "title": "原始注意力",
        "desc": "全部关注",
        "attention": original,
        "color": "tab:green"
    },
    {
        "title": "滑动窗口注意力",
        "desc": "W=3, i=5 → 关注位置 4,5,6",
        "attention": window,
        "color": "tab:blue"
    },
    {
        "title": "稀疏注意力",
        "desc": "W=2 + S=3, i=5 → 关注位置 1,4,5,6,8,10",
        "attention": sparse,
        "color": "tab:orange"
    }
]

# 创建子图
fig, axes = plt.subplots(3, 1, figsize=(12, 6))
fig.suptitle("注意力机制对比（序列长度=10）", fontsize=16, weight='bold')

for ax, row in zip(axes, rows):
    # 绘制每个位置
    for j, pos in enumerate(positions):
        is_attend = row["attention"][j]
        symbol = '●' if is_attend else '□'
        color = row["color"] if is_attend else 'lightgray'
        ax.text(pos, 0.5, symbol, fontsize=24, ha='center', va='center', color=color)
        ax.text(pos, 0.2, f'[{pos}]', fontsize=9, ha='center', va='center', color='black')
    
    # 设置坐标轴
    ax.set_xlim(0.5, n + 0.5)
    ax.set_ylim(0, 1)
    ax.axis('off')
    
    # 添加左侧标题（竖排）
    ax.text(-0.05, 0.5, row["title"], transform=ax.transAxes,
            rotation=90, va='center', ha='center',
            fontsize=12, weight='bold', color=row["color"])
    
    # 添加顶部说明文字
    ax.text(0.5, 0.92, row["desc"], transform=ax.transAxes,
            fontsize=11, ha='center', va='top',
            bbox=dict(boxstyle="round,pad=0.3", facecolor="whitesmoke"))

# 调整布局
plt.tight_layout(rect=[0.05, 0, 1, 0.95])
plt.show()

import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 设置序列长度范围
n_values = np.linspace(1000, 10000, 500)

# 计算不同模型下的计算量（相对值）
original_attention = n_values ** 2 / (1000 ** 2)  # 归一化处理以便比较
sliding_window = n_values * 0.001  # 假设比例因子使得其在图中显示合理
sparse_attention = n_values * 0.003  # 稀疏注意力斜率稍高于滑动窗口

# 创建图表
plt.figure(figsize=(10, 6))
plt.tight_layout(pad=0)
plt.subplots_adjust(left=0.08, right=0.98, top=0.95, bottom=0.1)

# 绘制三条曲线
plt.plot(n_values, original_attention, label='原始注意力 ( $ O(n^2) $ )', color='red')
plt.plot(n_values, sliding_window, label='滑动窗口 ( $ O(n) $ )', linestyle='--', color='blue')
plt.plot(n_values, sparse_attention, label='稀疏注意力 ( $ O(n) $ )', linestyle='-.', color='green')

# 添加标题和标签
plt.title('计算量随序列长度的变化曲线')
plt.xlabel('序列长度 n')
plt.ylabel('计算量（相对值）')

# 显示图例
plt.legend()

# 展示图表
plt.grid(True)
plt.show()