Croppable Knowledge Graph Embedding¶
会议: ACL 2025
arXiv: 2407.02779
代码: 基于 OpenKE 框架实现
领域: LLM / NLP
关键词: knowledge graph embedding, croppable model, mutual learning, multi-dimensional, parameter sharing
一句话总结¶
提出 MED 框架训练"可裁剪"知识图谱嵌入——一次训练同时优化 64 个不同维度的子模型（共享嵌入前缀），通过互学习、进化改进和动态损失权重，各维度子模型直接裁剪使用即超越独立训练和蒸馏方法，训练速度快 10 倍。
研究背景与动机¶
- 领域现状：KGE 将实体/关系映射到向量空间用于链路预测等任务。不同部署场景对嵌入维度需求不同——服务器可用 500 维，移动设备只能用 10 维。
- 现有痛点：每需要一个新维度就要从头训练一个新模型，或从高维教师模型做知识蒸馏。64 种维度需求 = 64 次独立训练，成本极高（如 WN18RR 上直接训练 64 个模型需 141 小时）。
- 核心矛盾：高维/低维模型的需求同时存在，但独立训练不共享信息、蒸馏方法（DualDE）虽然通过高维指导低维但训练时间更长（240 小时）。
- 本文要解决什么：一次训练得到可裁剪模型，截取嵌入前 \(d_i\) 维即可直接作为 \(d_i\) 维 KGE 使用。
- 切入角度：借鉴计算机视觉领域的 slimmable networks 思想，让不同维度的子模型共享参数前缀，并通过互学习和课程式训练使各维度都达到最优。
- 核心 idea 一句话：嵌入向量的前 \(d_1\) 维是最小子模型，前 \(d_2\) 维是次小子模型...前 \(d_n\) 维是最大子模型，一套参数通过 MED 三机制联合优化。
方法详解¶
整体框架¶
\(n=64\) 个子模型 \(M_1,...,M_n\) 共享嵌入参数前缀（\(d_1=10, d_2=20,...,d_{64}=640\)）。每个三元组 \((h,r,t)\) 在子模型 \(M_i\) 中用嵌入前 \(d_i\) 维计算评分 \(s^i_{(h,r,t)}\)。训练时同时优化所有子模型，通过三种机制确保各维度性能。
关键设计¶
- 互学习机制 (Mutual Learning):
- 做什么：让相邻维度子模型互为师生，双向传递知识
- 核心思路：\(L_{ML}^{i-1,i} = \sum_{(h,r,t)} d_\delta(s^{i-1}_{(h,r,t)}, s^i_{(h,r,t)})\)，使用 Huber loss (\(\delta=1\)) 衡量相邻子模型评分差异
-
设计动机：低维模型从高维学到更好的评分模式（知识蒸馏效果），高维模型从低维学到保持低维特征的稳定性（防止高维覆盖低维时性能退化）
-
进化改进机制 (Evolutionary Improvement):
- 做什么：让高维子模型聚焦于低维子模型预测不好的难样本
- 核心思路：对正样本，低维评分低的三元组获得更高权重 \(pos^i \propto \exp(w_1 / s^{i-1})\)；对负样本，低维评分高的（漏网之鱼）获得更高权重 \(neg^i \propto \exp(w_2 \cdot s^{i-1})\)
-
设计动机：高维子模型不应简单重复低维已做好的工作，而应专注于低维无法解决的难三元组，实现"进化式"能力提升
-
动态损失权重 (Dynamic Loss Weight):
- 做什么：自适应平衡不同维度子模型的训练重点
- 核心思路：总损失 \(L = \sum_{i=2}^n L_{ML}^{i-1,i} + \sum_{i=1}^n \exp(w_3 \cdot d_i / d_n) \cdot L_{EI}^i\)
- 设计动机：低维模型更需要从高维软标签学习（互学习主导），高维模型更需要关注硬标签和难样本（进化改进主导）
损失函数 / 训练策略¶
Adam 优化器 + 线性学习率衰减，batch size 1024，每个正样本生成 64 个负样本，最多 3000 轮 + 早停。\(w_1, w_2, w_3\) 为可学习参数（初始化为 1）。适用于任何有三元组评分函数的 KGE 方法。
实验关键数据¶
主实验 (RotatE on WN18RR)¶
| 维度 | MED MRR | MED H@10 | DualDE MRR | DualDE H@10 | DT MRR | DT H@10 |
|---|---|---|---|---|---|---|
| 10d | 0.324 | 0.469 | 0.179 | 0.440 | 0.131 | 0.303 |
| 40d | 0.466 | 0.561 | 0.441 | 0.525 | 0.447 | 0.528 |
| 160d | 0.471 | 0.574 | 0.465 | 0.560 | 0.470 | 0.575 |
| 640d | 0.476 | 0.574 | 0.476 | 0.576 | 0.508 | 0.612 |
消融实验 (训练效率)¶
| 方法 | WN18RR 训练时间 | 相对 MED | FB15K237 训练时间 | 相对 MED |
|---|---|---|---|---|
| MED | 12.7h | 1.0× | 35.5h | 1.0× |
| 直接训练 64 个 | 141.0h | 11.1× | 381.0h | 10.7× |
| DualDE | 240.0h | 18.9× | - | - |
| BKD | 163.0h | 12.8× | - | - |
关键发现¶
- 低维（10d）MED 的 MRR 比 DualDE 高 81%（RotatE WN18RR），因为 DualDE 蒸馏过程信息损失大
- 高维（640d）MED 比独立训练略低（0.476 vs 0.508），符合多目标优化的 trade-off 预期
- 训练速度比直接训练快 10 倍，比 DualDE 快 19 倍
- 实际电商场景（淘宝 SKG，7M 用户/50.7M 三元组）：移动端 10d 产品推荐 ndcg@5=0.422 > DT 0.344 > DualDE 0.404
- 扩展到 BERT：BERT-MED 在 GLUE 上比 BERT-HAT 平均高 16-22%，证明框架通用性
- 子模型数量从 64 减到 4 时，10d MRR 仅从 0.324 降到 0.319，但训练时间从 12.7h 降到 3.3h
亮点与洞察¶
- "可裁剪"的概念非常实用：一次训练，按需截取维度直接用，完美匹配异构设备部署需求。这个思路可迁移到任何嵌入式模型。
- 进化改进机制的"难样本聚焦"设计很巧妙：高维模型不重复低维已解决的问题，而是专攻难题，使维度增加真正带来能力增量。
- 互学习的双向性是关键：不只是高→低的蒸馏，低→高的反馈帮助高维模型保持低维特征稳定，这是 MED 比纯蒸馏方法强的原因。
局限性 / 可改进方向¶
- 实验仅覆盖 4 种 KGE 方法（TransE, SimplE, RotatE, PairRE），更强的方法（如 CompGCN）未验证
- 最大维度（640d）的性能略低于独立训练，在精度要求极高的场景可能不够
- 子模型数量增加时训练开销线性增长，极大量维度需求时仍有扩展性问题
- BERT 扩展实验仅做了 GLUE，未在更多 NLP 任务上验证
相关工作与启发¶
- vs DualDE (Zhu et al. 2023)：DualDE 用双向知识蒸馏，但需要先训练教师再蒸馏学生，训练时间是 MED 的 19 倍且低维性能差
- vs Slimmable Networks (Yu et al. 2019)：借鉴了 CV 中的可裁剪网络思想，但针对 KGE 的三元组评分做了特化设计
- 这个"可裁剪训练"范式可迁移到词嵌入、句子嵌入、甚至 LLM embedding layer 的弹性部署
评分¶
- 新颖性: ⭐⭐⭐⭐ "可裁剪 KGE"概念新颖，MED 三机制设计合理
- 实验充分度: ⭐⭐⭐⭐⭐ 4 种 KGE × 2 数据集 + 效率分析 + 工业应用 + BERT 扩展
- 写作质量: ⭐⭐⭐⭐ 动机清晰，实验详尽
- 价值: ⭐⭐⭐⭐ 解决实际部署痛点，工程价值高
-
仅在相邻子模型间蒸馏——维度差距小，蒸馏效果更好。
-
进化改进机制 (Evolutionary Improvement):
- 高维子模型需掌握低维子模型无法正确预测的三元组。
- 正例：\(M_{i-1}\) 评分低→ \(M_i\) 优化权重高（\(pos \propto \exp(w_1 / s^{i-1})\)）。
- 负例：\(M_{i-1}\) 评分高→ \(M_i\) 优化权重高（\(neg \propto \exp(w_2 \cdot s^{i-1})\)）。
-
使高维模型聚焦低维模型的"盲区"。
-
动态损失权重 (Dynamic Loss Weight):
- 低维子模型更依赖互学习（软标签），高维更依赖进化改进（硬标签）。
- 损失：\(L = \sum L_{ML} + \sum \exp(w_3 \cdot d_i / d_n) \cdot L_{EI}^i\)，高维子模型的进化改进权重指数增长。
实验关键数据¶
| 评估 | 结果 | 说明 |
|---|---|---|
| 4 KGE × 4 数据集 | MED 各维度子模型均表现良好 | 一次训练多维度直接用 |
| vs DualDE (SOTA蒸馏) | 低维子模型超越 | 且训练效率远高 |
| vs 独立训练 | 性能相当或更优 | 训练成本大幅降低 |
| 真实场景大规模 KG | 有效 | 实用性验证 |
| BERT + GLUE 扩展 | 有效 | 非 KGE 也适用 |
关键发现¶
- 三个机制缺一不可——互学习提升低维，进化改进聚焦盲区，动态权重平衡全局。
- 维度差距小的邻居蒸馏比隔远蒸馏效果好，验证了只在相邻子模型间学习的设计。
亮点与洞察¶
- "可裁剪"概念极其实用：一次训练 → 按需裁剪 → 直接部署，大幅降低多设备 KGE 的运维成本。
- 进化改进的"聚焦盲区"思路巧妙：高维不是简单重复低维已会的，而是专攻低维不会的。
- BERT 扩展验证说明 MED 框架不限于 KGE，具有通用性。
局限性 / 可改进方向¶
- 子模型数量 \(n\) 和各维度 \(d_i\) 需预先设定，无法动态调整。
- 仅在中等规模 KG 上验证，超大规模 KG 的效果有待测试。
评分¶
- 新颖性: ⭐⭐⭐⭐ "可裁剪"KGE 概念新颖，三机制设计完整
- 实验充分度: ⭐⭐⭐⭐⭐ 4 模型 × 4 数据集 + 真实场景 + BERT 扩展
- 写作质量: ⭐⭐⭐⭐ 公式推导清晰
- 价值: ⭐⭐⭐⭐ 大幅降低多维度 KGE 的训练和部署成本