对比加入和不加入缩放softmax的值的分布

gistfile1.txt

import numpy as np
import matplotlib.pyplot as plt

plt.rcParams["font.sans-serif"] = ["Songti SC"]
plt.rcParams["axes.unicode_minus"] = False

def softmax(x):
    """计算softmax并处理数值稳定性"""
    e_x = np.exp(x - np.max(x))  # 减去最大值防止指数溢出
    return e_x / e_x.sum()

# 参数设置
d_k = 64  # 向量维度（模拟实际场景）
np.random.seed(42)  # 固定随机种子保证结果可复现

# 生成模拟的注意力分数（未缩放）
raw_scores = np.random.randn(50) * np.sqrt(d_k)  # 假设q和k的点积方差与d_k成正比
scaled_scores = raw_scores / np.sqrt(d_k)  # 缩放后的分数

# 计算softmax结果
softmax_raw = softmax(raw_scores)
softmax_scaled = softmax(scaled_scores)

# 绘制对比图
plt.figure(figsize=(10, 5))
x = np.arange(len(raw_scores))

plt.bar(x - 0.2, softmax_raw, width=0.4, label='未缩放 (Raw Scores)', alpha=0.7, color='red')
plt.bar(x + 0.2, softmax_scaled, width=0.4, label='缩放后 (Scaled Scores)', alpha=0.7, color='blue')

plt.title('Softmax分布对比：缩放 vs 未缩放 (d_k=64)')
plt.xlabel('注意力位置索引')
plt.ylabel('Softmax概率')
plt.xticks(x)
plt.legend()
plt.grid(linestyle='--', alpha=0.5)
plt.savefig(f"./images/softmax-stat.png")
plt.show()