4. 关系蒸馏：样本间关系保持、基于距离与角度的蒸馏方法

前面我们聊了特征层怎么对齐，那是让单个样本的中间表示尽量像老师。但有个问题——样本之间的关系，我们还没管。

你想想看，一个学生模型如果只学单样本的特征，它可能记住了「猫长这样」，但不知道「猫和狗比，谁更像老虎」。这种样本间的相对位置，才是知识迁移里更高级的东西。

关系蒸馏（Relational Knowledge Distillation, RKD）就是干这个的。它的核心思想很简单：老师模型内部，样本 A 和样本 B 的距离、角度、拓扑关系，学生模型也要学过来。

4.1 为什么需要关系蒸馏？

我在一个视觉检索项目里踩过坑。当时用特征蒸馏训了一个轻量模型，单样本分类准确率还行，但一做检索（比如找相似图），结果就乱套了。学生模型把两张不同品种的猫图判成不相似，把猫和狗判成相似——因为它没学到样本间的相对距离。

关系蒸馏的价值就在这里：

• 保持语义结构：老师模型在特征空间里，同类样本聚在一起，不同类分开。学生也要保持这个结构。
• 提升泛化能力：学关系比学单点特征更鲁棒。即使某个样本特征有偏差，关系结构还能兜底。
• 适合度量学习任务：人脸识别、行人重识别、图像检索，这些任务本质就是比关系。

4.2 关系蒸馏的核心框架

关系蒸馏的通用框架分两步：

1. 定义关系函数：从特征空间中提取样本对（或样本三元组）的关系度量，比如距离、角度、互信息。
2. 最小化关系差异：让学生模型的关系输出，尽量接近老师模型的关系输出。

用公式表达就是：

L_RKD = Σ_{(i,j) ∈ P} D( ψ_T(f_T(x_i), f_T(x_j)), ψ_S(f_S(x_i), f_S(x_j)) )

其中 ψ 是关系函数，D 是差异度量（比如 L2 损失），P 是样本对集合。

下面这张图展示了关系蒸馏的整体流程：

4.3 基于距离的关系蒸馏

这是最直观的方法。我们计算样本对在特征空间中的欧氏距离或余弦距离，然后让学生模型的距离分布匹配老师模型。

4.3.1 距离蒸馏损失

给定一个 mini-batch 里的样本对 (i, j)，定义距离关系：

d_T(i,j) = || f_T(x_i) - f_T(x_j) ||₂
d_S(i,j) = || f_S(x_i) - f_S(x_j) ||₂

L_dist = (1 / |P|) * Σ_{(i,j) ∈ P} (d_T(i,j) - d_S(i,j))²

这里 P 是所有样本对的集合。如果 batch size 是 N，那就有 N×(N-1)/2 个对。

4.3.2 距离蒸馏的变体

方法	关系定义	特点
欧氏距离蒸馏	L2 距离	简单直接，对尺度敏感
余弦距离蒸馏	1 - cos(θ)	对尺度不敏感，关注方向
归一化距离蒸馏	先 L2 归一化再算 L2	综合两者优点，我常用这个
相对距离蒸馏	d(i,j) / mean(d)	关注相对远近，而非绝对值

4.4 基于角度的关系蒸馏

距离只刻画了样本对的「远近」，但没刻画「方向」。举个例子：三个样本 A、B、C，A 和 B 的距离等于 A 和 C 的距离，但 B 和 C 可能一个在 A 的左边，一个在 A 的右边。角度关系能捕捉这种差异。

4.4.1 角度蒸馏损失

角度蒸馏通常用三元组 (i, j, k)，计算以 i 为顶点的夹角：

cos_θ_T(i,j,k) = ( (f_T(x_j) - f_T(x_i)) · (f_T(x_k) - f_T(x_i)) ) / ( ||f_T(x_j) - f_T(x_i)||₂ · ||f_T(x_k) - f_T(x_i)||₂ )

L_angle = (1 / |T|) * Σ_{(i,j,k) ∈ T} (cos_θ_T(i,j,k) - cos_θ_S(i,j,k))²

其中 T 是所有三元组的集合。如果 batch size 是 N，三元组数量是 N×(N-1)×(N-2)/2，计算量比较大。

4.4.2 距离 vs 角度：怎么选？

维度	距离蒸馏	角度蒸馏
计算复杂度	O(N²)	O(N³)
信息粒度	粗粒度（远近）	细粒度（方位）
对特征尺度	敏感	不敏感（余弦自带归一化）
适用场景	分类、检索	细粒度检索、聚类
训练稳定性	高	中等（三元组选择影响大）

4.5 关系蒸馏的进阶技巧

4.5.1 多尺度关系融合

单一的关系函数可能不够。我习惯把距离和角度结合起来：

L_RKD = α * L_dist + β * L_angle

α 和 β 是平衡系数。我一般从 α=1.0, β=0.1 开始调，因为距离蒸馏更稳定，角度蒸馏作为辅助。

4.5.2 关系蒸馏 + 特征蒸馏联合训练

别把关系蒸馏和特征蒸馏对立起来。它们互补：

L_total = L_CE + λ1 * L_feature + λ2 * L_RKD

L_CE 是分类损失，L_feature 是特征蒸馏损失（比如前面章节讲的），L_RKD 是关系蒸馏损失。

4.5.3 关系蒸馏的采样策略

全量样本对计算量太大。我常用两种采样策略：

• 随机采样：每个 batch 随机选一部分对，简单但可能漏掉关键关系。
• 困难对采样：选那些老师模型和学生模型关系差异最大的对，针对性优化。收敛更快，但需要额外计算差异。

我个人更倾向困难对采样。虽然每步计算多一点，但总训练步数能减少 30% 左右。

4.6 实战：关系蒸馏的 PyTorch 实现

下面给一个距离蒸馏的参考实现。代码不长，核心就几十行：

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

class DistanceDistillationLoss(nn.Module):
    def __init__(self, normalize=True, margin=None):
        super().__init__()
        self.normalize = normalize
        self.margin = margin  # 可选：只蒸馏距离大于阈值的对

    def forward(self, feat_s, feat_t):
        # feat_s, feat_t: [batch_size, dim]
        if self.normalize:
            feat_s = F.normalize(feat_s, p=2, dim=1)
            feat_t = F.normalize(feat_t, p=2, dim=1)

        # 计算距离矩阵
        dist_s = torch.cdist(feat_s, feat_s, p=2)  # [N, N]
        dist_t = torch.cdist(feat_t, feat_t, p=2)

        # 取上三角（排除对角线）
        n = dist_s.size(0)
        mask = torch.triu(torch.ones(n, n, device=feat_s.device), diagonal=1).bool()
        d_s = dist_s[mask]
        d_t = dist_t[mask]

        # 可选：只蒸馏距离大于 margin 的对
        if self.margin is not None:
            keep = d_t > self.margin
            d_s = d_s[keep]
            d_t = d_t[keep]

        # MSE 损失
        loss = F.mse_loss(d_s, d_t)
        return loss

4.7 关系蒸馏的常见问题与调试

问题	可能原因	解决方案
关系损失不下降	特征尺度不匹配	加 L2 归一化，或调小学习率
分类准确率反而下降	关系损失权重太大	降低 λ2，或先单独训分类再联合
训练速度太慢	全量样本对计算	用随机采样或困难对采样
角度蒸馏不收敛	batch size 太小	增大 batch size 或用 memory bank
检索任务效果差	只用了距离蒸馏	加上角度蒸馏，或换相对距离