第 8 章 · TransE 与注意力代码实验¶

# TransE 训练：让正确三元组距离变小，让替换尾实体后的三元组距离变大。
entities = ["France", "Germany", "Italy", "Spain", "Portugal", "Paris", "Berlin", "Rome", "Madrid", "Lisbon"]
capitals = ["Paris", "Berlin", "Rome", "Madrid", "Lisbon"]
entity_to_id = {name: idx for idx, name in enumerate(entities)}
positive_triples = [
    ("France", "capital_of", "Paris"),
    ("Germany", "capital_of", "Berlin"),
    ("Italy", "capital_of", "Rome"),
    ("Spain", "capital_of", "Madrid"),
    ("Portugal", "capital_of", "Lisbon"),
]

rng = np.random.default_rng(8)
entity_vec = rng.normal(0, 0.18, size=(len(entities), 2))
relation_vec = rng.normal(0, 0.18, size=2)
margin = 0.45
lr = 0.035
loss_rows = []

def triple_distance(head, tail):
    h = entity_vec[entity_to_id[head]]
    t = entity_vec[entity_to_id[tail]]
    return float(np.linalg.norm(h + relation_vec - t))

for epoch in range(1, 501):
    total_loss = 0.0
    rng.shuffle(positive_triples)
    for head, _, true_tail in positive_triples:
        negative_tail = rng.choice([tail for tail in capitals if tail != true_tail])
        h_id = entity_to_id[head]
        t_id = entity_to_id[true_tail]
        n_id = entity_to_id[negative_tail]

        pos_vec = entity_vec[h_id] + relation_vec - entity_vec[t_id]
        neg_vec = entity_vec[h_id] + relation_vec - entity_vec[n_id]
        pos_score = float(np.dot(pos_vec, pos_vec))
        neg_score = float(np.dot(neg_vec, neg_vec))
        loss = max(0.0, margin + pos_score - neg_score)
        total_loss += loss
        if loss == 0:
            continue

        grad_pos = 2 * pos_vec
        grad_neg = 2 * neg_vec
        entity_vec[h_id] -= lr * (grad_pos - grad_neg)
        relation_vec -= lr * (grad_pos - grad_neg)
        entity_vec[t_id] += lr * grad_pos
        entity_vec[n_id] -= lr * grad_neg

    if epoch in {1, 20, 50, 100, 200, 500}:
        pos_mean = np.mean([triple_distance(h, t) for h, _, t in positive_triples])
        neg_examples = [(h, rng.choice([tail for tail in capitals if tail != t])) for h, _, t in positive_triples]
        neg_mean = np.mean([triple_distance(h, t) for h, t in neg_examples])
        loss_rows.append({"训练轮次": epoch, "平均损失": total_loss / len(positive_triples), "正例平均距离": pos_mean, "负例平均距离": neg_mean})

transe_loss_df = pd.DataFrame(loss_rows)
eval_rows = []
for head, relation, true_tail in positive_triples:
    candidates = sorted(
        [{"头实体": head, "候选尾实体": tail, "距离": triple_distance(head, tail), "是否正确": tail == true_tail} for tail in capitals],
        key=lambda row: row["距离"],
    )
    eval_rows.extend(candidates[:3])

transe_df = pd.DataFrame(eval_rows)
display(transe_loss_df.round(3))
display(transe_df.round(3))

# 绘制训练后的实体位置、关系平移和损失变化。
fig, axes = plt.subplots(1, 2, figsize=(11.0, 4.8))
ax = axes[0]
for name in entities:
    vec = entity_vec[entity_to_id[name]]
    is_capital = name in capitals
    ax.scatter(vec[0], vec[1], s=150, color="#fff7ed" if is_capital else "#eff6ff", edgecolor="#f97316" if is_capital else "#2563eb", linewidth=1.5)
    ax.text(vec[0] + 0.015, vec[1] + 0.015, name, color="#0f172a")

for head, _, tail in positive_triples[:3]:
    h = entity_vec[entity_to_id[head]]
    t = entity_vec[entity_to_id[tail]]
    predicted = h + relation_vec
    ax.arrow(h[0], h[1], relation_vec[0], relation_vec[1], width=0.004, head_width=0.035, color="#2563eb", alpha=0.55, length_includes_head=True)
    ax.plot([predicted[0], t[0]], [predicted[1], t[1]], color="#f97316", linewidth=1.2, linestyle="--")

ax.set_title("训练后：h + r 靠近正确首都", loc="left", fontweight="bold", color="#0f172a")
ax.set_xlabel("维度1")
ax.set_ylabel("维度2")
ax.grid(True, color="#e2e8f0", linewidth=0.8)

axes[1].plot(transe_loss_df["训练轮次"], transe_loss_df["正例平均距离"], marker="o", color="#2563eb", linewidth=2.0, label="正例距离")
axes[1].plot(transe_loss_df["训练轮次"], transe_loss_df["负例平均距离"], marker="o", color="#f97316", linewidth=2.0, label="负例距离")
axes[1].set_title("训练过程", loc="left", fontweight="bold", color="#0f172a")
axes[1].set_xlabel("训练轮次")
axes[1].set_ylabel("平均距离")
axes[1].grid(True, color="#e2e8f0", linewidth=0.8)
axes[1].legend()
plt.tight_layout()
plt.show()

# 用一组可解释的查询/键特征，让同一句话的注意力更容易读。
labels = ["这", "只", "猫", "坐", "在", "垫子"]
features = ["subject_noun", "verb", "prep", "place_noun", "det"]

K = pd.DataFrame([
    [0, 0, 0, 0, 1],
    [1, 0, 0, 0, 0],
    [0, 1, 0, 0, 0],
    [0, 0, 1, 0, 0],
    [0, 0, 0, 0, 1],
    [0, 0, 0, 1, 0],
], index=labels, columns=features)

Q = pd.DataFrame([
    [8.0, 0.0, 0.0, 0.0, 0.0],
    [0.0, 8.0, 0.0, 0.0, 0.0],
    [6.5, 0.0, 0.0, 6.5, 0.0],
    [0.0, 0.0, 0.0, 8.0, 0.0],
    [0.0, 0.0, 0.0, 8.0, 0.0],
    [0.0, 0.0, 8.0, 0.0, 0.0],
], index=labels, columns=features)

V = np.array([
    [0.20, 0.10],
    [0.85, 0.20],
    [0.45, 0.80],
    [0.20, 0.65],
    [0.20, 0.10],
    [0.70, 0.55],
])

scores = Q.to_numpy() @ K.to_numpy().T / math.sqrt(len(features))
weights = np.exp(scores - scores.max(axis=1, keepdims=True))
weights = weights / weights.sum(axis=1, keepdims=True)
attn_output = weights @ V

weight_df = pd.DataFrame(np.round(weights, 3), index=labels, columns=labels)
focus_rows = []
for i, label in enumerate(labels):
    kept = [j for j in range(len(labels)) if weights[i, j] >= 0.20]
    focus_rows.append({
        "当前词": label,
        "主要关注": ", ".join(labels[j] for j in kept),
        "权重": ", ".join(f"{weights[i, j]:.2f}" for j in kept),
    })

display(pd.DataFrame(focus_rows))
display(weight_df)
display(pd.DataFrame(np.round(attn_output, 3), index=labels, columns=["输出1", "输出2"]))

# 画出权重热力图，并把每一行的最高权重直接标出来。
fig, (ax, link_ax) = plt.subplots(
    1,
    2,
    figsize=(10.8, 5.2),
    gridspec_kw={"width_ratios": [1.25, 0.9]},
)
im = ax.imshow(weights, cmap="YlOrRd", vmin=0, vmax=1)
ax.set_xticks(range(len(labels)), labels, rotation=30, ha="right")
ax.set_yticks(range(len(labels)), labels)
ax.set_xlabel("被关注的词（键）")
ax.set_ylabel("正在读的词（查询）")

for i in range(len(labels)):
    max_weight = weights[i].max()
    for j in range(len(labels)):
        value = weights[i, j]
        color = "#ffffff" if value > 0.45 else "#0f172a"
        weight = "bold" if value == max_weight else "normal"
        ax.text(j, i, f"{value:.2f}", ha="center", va="center", color=color, fontweight=weight)
        if value == max_weight:
            ax.add_patch(plt.Rectangle((j - 0.5, i - 0.5), 1, 1, fill=False, edgecolor="#0f172a", linewidth=2.4))

ax.set_title("注意力权重热力图", loc="left", fontsize=14, fontweight="bold", color="#0f172a")
fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

link_ax.set_xlim(0, 1)
link_ax.set_ylim(-0.7, len(labels) - 0.25)
link_ax.axis("off")
link_ax.text(0.02, len(labels) - 0.15, "每行最高关注", fontsize=13, fontweight="bold", color="#0f172a")
link_ax.text(0.08, len(labels) - 0.65, "当前词", color="#64748b", fontweight="bold")
link_ax.text(0.68, len(labels) - 0.65, "被关注词", color="#64748b", fontweight="bold")

for i in range(len(labels)):
    y = len(labels) - 1 - i
    max_weight = weights[i].max()
    top_indices = [j for j in range(len(labels)) if np.isclose(weights[i, j], max_weight)]
    target_text = " / ".join(labels[j] for j in top_indices)
    value_text = " / ".join(f"{weights[i, j]:.2f}" for j in top_indices)
    link_ax.text(0.08, y, labels[i], ha="left", va="center", color="#0f172a")
    link_ax.text(0.68, y, target_text, ha="left", va="center", color="#0f172a", fontweight="bold")
    link_ax.annotate(
        "",
        xy=(0.65, y),
        xytext=(0.35, y),
        arrowprops={"arrowstyle": "->", "lw": 1.2 + 4.0 * max_weight, "color": "#ea580c", "alpha": 0.95},
    )
    link_ax.text(0.50, y + 0.16, value_text, ha="center", va="center", color="#c2410c", fontsize=9)

plt.tight_layout()
plt.show()