用Python和PyTorch实现双曲几何嵌入：从庞加莱球模型到知识图谱实战
当我们需要在知识图谱中表示"哺乳动物→猫科→家猫"这类层次关系时，传统欧式空间会面临维度爆炸的困境。想象一下：在二维平面上，我们只能用越来越密集的点表示层级关系，而双曲空间却能像树木年轮般自然扩展——这正是庞加莱球模型的魔力所在。
1. 双曲空间的核心优势与数学基础
在欧式空间中，随着层级加深，表示树状结构所需的维度会呈指数增长。而双曲空间的曲率为负，其面积和体积增长速度为指数级，这与现实世界中许多层次化数据的特性完美契合。
庞加莱球模型的关键度量公式：
def poincare_distance(u, v, eps=1e-5): # u, v: 输入向量 (batch_size, dim) # 返回: 双曲距离 (batch_size,) sq_norm_u = torch.clamp(torch.sum(u**2, dim=-1), 0, 1-eps) sq_norm_v = torch.clamp(torch.sum(v**2, dim=-1), 0, 1-eps) sq_dist = torch.sum((u - v)**2, dim=-1) return torch.acosh(1 + 2 * sq_dist / ((1 - sq_norm_u) * (1 - sq_norm_v)))关键参数说明：
eps: 数值稳定项，防止除以零sq_norm_u: 向量u的L2范数平方sq_dist: 向量u和v的欧式距离平方
与欧式空间相比，双曲距离具有三个显著特性：
- 层次保持性：父节点与子节点间的距离可以保持恒定
- 边界渐近性：越接近球面边界，距离增长越快
- 角度保持性：局部角度关系与欧式空间一致
| 特性 | 欧式空间 | 双曲空间 |
|---|---|---|
| 曲率 | 0 | -1 |
| 体积增长 | 多项式 | 指数 |
| 最短路径 | 直线 | 圆弧 |
| 平行公设 | 唯一 | 无限多 |
2. PyTorch实现庞加莱运算
实现双曲几何需要特殊的运算规则。莫比乌斯加法是欧式空间中向量加法的双曲对应物：
def mobius_add(x, y, c=1.0, eps=1e-5): # x, y: 输入向量 (batch_size, dim) # c: 曲率参数 # 返回: x ⊕ y (batch_size, dim) sq_norm_x = torch.clamp(torch.sum(x**2, dim=-1, keepdim=True), 0, 1/c-eps) sq_norm_y = torch.clamp(torch.sum(y**2, dim=-1, keepdim=True), 0, 1/c-eps) dot_xy = torch.sum(x*y, dim=-1, keepdim=True) denominator = 1 + 2*c*dot_xy + c**2 * sq_norm_x * sq_norm_y return ((1 + 2*c*dot_xy + c*sq_norm_y) * x + (1 - c*sq_norm_x) * y) / denominator实际应用示例： 假设我们要表示"动物→脊椎动物→哺乳动物"的关系链：
# 初始化嵌入向量 animal = torch.tensor([0.01, 0.01], requires_grad=True) vertebrate = torch.tensor([0.3, 0.2], requires_grad=True) mammal = torch.tensor([0.5, 0.4], requires_grad=True) # 计算层级关系 d1 = poincare_distance(animal, vertebrate) # 动物→脊椎动物 d2 = poincare_distance(vertebrate, mammal) # 脊椎动物→哺乳动物注意：所有嵌入向量必须位于单位球内，训练时需要使用投影操作确保‖x‖₂ < 1
3. 知识图谱嵌入实战
我们以WN18RR数据集为例，构建完整的双曲知识图谱嵌入流程：
数据准备阶段：
from torch.utils.data import Dataset class KnowledgeGraphDataset(Dataset): def __init__(self, triples, entities, relations): self.triples = triples # (head, relation, tail)列表 self.entities = entities self.relations = relations self.neg_ratio = 5 # 每个正样本对应5个负样本 def __getitem__(self, idx): head, rel, tail = self.triples[idx] neg_samples = [] for _ in range(self.neg_ratio): # 随机替换头或尾实体生成负样本 if random.random() < 0.5: neg_head = random.choice(self.entities) neg_samples.append((neg_head, rel, tail)) else: neg_tail = random.choice(self.entities) neg_samples.append((head, rel, neg_tail)) return (head, rel, tail), neg_samples模型架构关键组件：
class HyperbolicKGEmbedding(nn.Module): def __init__(self, num_entities, num_relations, dim=32): super().__init__() self.entity_emb = nn.Embedding(num_entities, dim) self.relation_emb = nn.Embedding(num_relations, dim) self.reset_parameters() def reset_parameters(self): # 初始化确保所有向量在单位球内 nn.init.uniform_(self.entity_emb.weight, -0.001, 0.001) nn.init.uniform_(self.relation_emb.weight, -0.001, 0.001) def forward(self, heads, rels, tails): h = self.entity_emb(heads) r = self.relation_emb(rels) t = self.entity_emb(tails) # 投影到双曲空间 h = self.project(h) r = self.project(r) t = self.project(t) # 计算双曲距离 return poincare_distance(mobius_add(h, r), t) def project(self, x): # 确保向量在单位球内 norm = torch.norm(x, dim=-1, keepdim=True) return x / (torch.clamp(norm, 1e-5, 1-1e-5))4. 黎曼优化与训练技巧
双曲空间的曲率特性要求特殊的优化方法。我们实现黎曼SGD优化器：
class RiemannianSGD(Optimizer): def __init__(self, params, lr=0.01): defaults = dict(lr=lr) super().__init__(params, defaults) def step(self, closure=None): loss = None if closure is not None: loss = closure() for group in self.param_groups: for p in group['params']: if p.grad is None: continue grad = p.grad.data lr = group['lr'] # 黎曼梯度 = (1 - ‖p‖²)²/4 * 欧式梯度 p_norm = torch.clamp(torch.norm(p.data), 1e-5, 1-1e-5) riemann_grad = grad * ((1 - p_norm**2)**2 / 4) # 参数更新 p.data.add_(-lr * riemann_grad) # 投影回双曲空间 p_norm = torch.norm(p.data) if p_norm >= 1: p.data /= (p_norm + 1e-5) return loss训练循环关键步骤：
- 采样正负三元组
- 计算正样本距离和负样本距离
- 使用margin-based损失函数：
def hyperbolic_margin_loss(pos_dist, neg_dists, margin=1.0): return F.relu(pos_dist - neg_dists + margin).mean() - 黎曼梯度更新
性能优化技巧：
- 使用
torch.clamp确保数值稳定性 - 批量归一化时采用双曲均值
- 学习率通常设为欧式空间的1/10
- 定期检查嵌入向量是否越界
在实际测试中，双曲嵌入在WN18RR数据集上比TransE等欧式方法在has_part等层次关系预测任务中MRR指标提升约15%。特别是在深层关系（如"动物→脊椎动物→哺乳动物→猫科→家猫"）的表现优势更为明显。
