NRI神经关系推理：从图结构潜变量到轨迹预测的PyTorch实现

把NRI拆开看，它做的核心事情其实就一件：从观测数据里自动推断出物体之间的交互关系，再用推断出的关系去预测未来的运动轨迹。这个思路听起来不复杂，但它切中的恰好是图神经网络在实际任务里最隐蔽的一个痛点——几乎所有常见GNN方法都假设“图是给定的”，可现实里你拿到一堆漂浮粒子的坐标、一群行人的追踪框、几辆车的GPS轨迹，它们之间谁跟谁有交互、交互是什么类型，根本没人替你标注。

NRI（Neural Relational Inference，神经关系推理）最早出现在ICML 2018，作者里有Thomas Kipf（GCN的作者）和Max Welling，后续在物理系统建模、多智能体轨迹预测、因果发现这些方向被反复引用。它的核心贡献是把“图结构”本身定义成潜变量，让模型在一个端到端的框架里同时完成关系推理和动力学预测，而不是先用启发式方法把图算出来，再丢给下游GNN。这篇文章我会把模型原理拆成能看懂、能复现的程度，配套一份基于PyTorch的弹簧质点系统运动轨迹预测代码，你可以直接跑通并在此基础上做自己的实验。适合有机器学习基础、会用PyTorch、刷过GNN入门项目的读者，如果你对图卷积、注意力机制只是听说过但没动手写过，也完全可以跟着代码走一遍，难点我都会解释。

1. 为什么运动轨迹预测要先做“关系推理”

1.1 一个反直觉的观察：GNN的上限由图的输入质量决定

很多初次接触图神经网络的读者会有一种直觉：只要网络够深、特征够丰富，模型应该能从数据里“自己学会”物体之间谁重要谁不重要。这个直觉在部分场景下成立，但那是因为数据集里的图本身就是人工标注好的。一旦切换到真实的动态系统预测场景，问题就暴露了。

以最简单也最经典的弹簧质点系统为例。你观测到N个质点在二维平面里运动，每个时刻每个质点有四个数值：横坐标x、纵坐标y、速度vx、vy。它们之间有一部分存在弹簧连接，另一部分没有任何物理关联。假设系统没有重力、空气阻力，那对任何一个质点来说，它在下一时刻的加速度完全由“谁和它之间有弹簧”决定。如果你拿到了正确的连接矩阵A，预测任务就退化成一个确定性的力学积分问题，即使不用神经网络，用胡克定律一步步推也能推出非常准的轨迹。

但问题在于，观察者能看到的只有位置和速度，而A是不可见的。这时候如果照搬GNN标准流程，直接把所有节点两两连接成一个完全图，然后让消息传递自己去学权重，会发生什么？实验结果是：预测效果远不如按真实连接做GNN。原因是完全图引入了大量虚假的交互路径，消息传递会把不存在的连接也纳入节点更新，相当于给每个质点加了额外的“幻觉力”。你让模型学的是一个被严重污染的函数，就算网络容量再大，也很难自动过滤掉这些错误的边。

GAT这类注意力模型能缓解一部分问题，因为它能学到边的软权重。但GAT学到的权重是连续标量，不具备“边类型”的语义，也没有显式的稀疏约束，在需要理解系统内在结构（比如“哪些边是弹簧、哪些边是排斥力、哪些边根本不存在”）的场景下，它的可解释性和结构恢复能力都打折扣。NRI的出发点正是这里：与其把结构当特征去隐式学习，不如把结构当成需要推理的潜变量，让模型先回答“这个系统长什么样”，再基于它做预测。

1.2 把“关系”当作潜变量：NRI的核心思想一句话版

NRI把轨迹建模成两个阶段。第一阶段，编码器看一段历史观测（比如前9帧的位置速度），输出每一对节点之间“边类型”的后验分布。第二阶段，解码器拿着从分布里采样出的离散图结构，用消息传递网络（实际上是带类型的消息函数）对系统未来的演化做多步解码。

这里的潜变量z_ij是离散的，表示节点i和节点j之间边的类型。如果数据里只有“有弹簧”和“没弹簧”，那z_ij就是一个二分类变量；更复杂的物理系统里可能有弹簧、排斥力、阻尼、万有引力等多种交互类型，z_ij就变成一个K分类变量。NRI模型本身并不限定K的取值，你完全可以根据你的系统语义来定义边类型的数量，哪怕K等于4、5、8都可以。

这样设计有一个很大的好处：推理出来的z_ij不仅用来预测未来轨迹，它本身就是一种可解释的结构输出。比如测试一个多体系统，模型跑了几个epoch之后把边恢复准确率打到95%以上，那这张边矩阵就是你可以直接拿去分析的“系统接线图”。这种副产品在很多实际工程场景里比轨迹预测本身还有价值。

1.3 和常见替代方案的对比

为了让你更直观地理解NRI在方法谱系里处于什么位置，我把它和几类常见做法放在一起对比：

从表格里能看出来，前两类方法本质上都在回避“这个系统到底长什么样”的问题。NRI选择正面回答它，并且用离散潜变量和Gumbel-Softmax把结构推理和轨迹预测整合进同一个可端到端训练的框架。这种思路在“结构即语义”的物理系统和多智能体系统里特别占优势。

2. NRI模型架构逐模块拆解：从边类型潜变量到多步解码

2.1 编码器：从观测序列到边类型的后验分布

NRI的编码器其实就是一个共享的GNN，输入是长度为T的历史轨迹，输出是每一对节点之间属于每种边类型的logits。它的工作流程可以分为两步。

第一步，把每个节点在时间维度的观测压缩成一个特征向量。常用的做法是用一个1D卷积或者多层感知机对节点的历史状态做编码，然后做时间维度的平均池化或取最后时刻，得到h_i^0。这一步处理的是序列信息，让编码器不仅能感知当前帧的位置速度，还能从中提取出运动模式的局部特征。

第二步，在节点层面上做几轮消息传递。消息传递的规则可以这么写：

h_i^{l+1} = f_emb(h_i^l) + Σ_{j ≠ i} f_msg(h_i^l, h_j^l)

其中f_emb和f_msg都是MLP。f_msg以发送者节点和接收者节点的特征拼接为输入，输出一条“消息”；所有邻居的消息求和后更新接收者的表示。这里没有用到任何图的先验信息，因为编码器要处理的就是一个完全图，所有节点两两之间都可能是候选边。消息传递在这里的作用是让节点的表示不仅包含自身运动信息，还包含它和所有邻居交互的聚合信息。

消息传递结束后，编码器对每一对节点(i,j)，把h_i和h_j拼接起来，通过一个线性层输出K维的logits：

logits_ij = W_out [h_i || h_j]

这里logits_ij的语义就是“在已知观测x的情况下，节点i和j之间边类型属于第k类的概率对数值”，所以它本质上建模的是后验分布q_phi(z_ij | x)。有一点需要注意，在物理系统中边通常是无向的，因此logits_ij和logits_ji会做对称化处理，比如相加或取平均，以保证推理出的图结构满足对称性。

2.2 Gumbel-Softmax：离散潜变量怎么反向传播

到了这一步，你手上有了logits_ij，接下来要从这个分布里采出一个离散的边类型z_ij作为解码器的输入。但是问题来了：离散采样是不可导的，梯度没办法从解码器的损失函数回流到编码器。你总不能把解码器的loss通过一个“掷骰子”的操作反向传播到logits上。

解决这个问题的经典手段是Gumbel-Softmax重参数化。简单理解，它做了一个“软采样”：把离散的one-hot向量换成连续的、温度可控的近似分布。Gumbel-Softmax采样过程可以直观看出它做了什么：

z ≈ softmax((logits + g) / tau)

其中g是从Gumbel分布采样的噪声，tau是温度参数。当tau趋近0时，这个softmax的输出会越来越接近一个one-hot向量，近似程度越高，但梯度方差也越大；当tau比较大时，输出更平滑，梯度更稳定，但离真实的离散采样更远。实际操作中，NRI论文里通常固定使用较大的温度（比如tau=1）训练编码器，让梯度稳定传播，测试阶段再用argmax得到真正离散的z。我的经验是，简单的固定温度训练效果就已经比很多人预想的好，温度退火不是必须的，这点后面在调参心得里再展开说。

2.3 解码器：基于推断结构的多步轨迹预测

解码器是NRI里决定预测上限的部分，也是和普通GNN最不一样的地方。它的核心设计是“按边类型分配独立的消息函数”。

拿到采样出的z_ij（已经是one-hot向量），解码器在每一步做这样几件事：

第一，根据每个节点的当前隐状态h_i^{t}，和每个邻居j的隐状态h_j^{t}，拼接后送入消息函数。关键在于，消息函数是按边类型分组的：每种边类型都有一个独立的MLP。比如第k种边类型对应一个MLP f_msg_k，那么对于每一条边(i,j)，只有当z_ij指示的类型是k时，这条边才会使用f_msg_k来计算消息。由于z_ij是one-hot向量，等价于把所有边类型对应的消息按权重加和，数学上可以写成：

msg_ij = Σ_k z_ij^k · f_msg_k(h_i, h_j)

第二，每个节点聚合所有邻居发来的消息，用GRU更新自己的隐状态：

h_i^{t+1} = GRU_emb(Σ_{j≠i} msg_ij, h_i^{t})

第三，从新的隐状态里解码出速度增量，再按运动学规律积分更新位置和速度。这里的做法是：把当前的位置和速度拼接起来，通过一个线性层得到这一步的预测输出，然后对速度做增量更新，再把位置按速度积分：

v^{t+1} = v^{t} + Δv^t x^{t+1} = x^{t} + v^{t+1} · dt

为什么要用GRU而不是普通的前馈网络？因为轨迹预测是典型的多步自回归过程，上一步的隐状态里包含了系统到目前为止的全部动力学信息。GRU的循环结构让信息能沿着时间轴传播，同时比LSTM参数更少，在后期的长轨迹预测里能有效降低过拟合风险。

这样一个按边类型独立消息函数的设计带来的效果是：不同类型的交互通过不同的函数来建模，模型不需要用一个“万能MLP”同时拟合弹簧力、排斥力、摩擦力这些千差万别的动力学，而是让每种交互自己学一套消息机制。这就是NRI能在复杂多体系统上恢复出结构并做出准确预测的根本原因。

2.4 训练目标：ELBO损失和推理阶段的差异

NRI的训练目标是最大化观测轨迹的对数似然的下界（ELBO）。直观理解分成两项：

第一项是重建损失，要求解码器在给定推断结构的情况下，尽可能准确地预测出未来每一帧的位置和速度。在具体实现里，这一项通常写成高斯负对数似然，等价于简化版的均方误差损失。

第二项是KL散度，要求编码器推断出来的结构分布不要和先验分布偏离太多。NRI论文用的是均匀先验，也就是没有任何外部信息时，所有边类型出现的概率应该接近等可能。不过需要注意的是，如果你的系统本身边类型就有很强的先验不平衡（比如99%的边都是不存在的），KL权重需要适当调低，否则模型会被先验拖住，不敢大胆预测“有边”。

训练阶段和推理阶段有一个微妙但很重要的差异。训练时z_ij是从Gumbel-Softmax分布里软采样出来的，所以解码器拿到的是“概率意义上的软图”；推理时我们期望的是离散的、可直接解释的边类型，所以用argmax把logits变成one-hot。如果你在推理阶段仍然使用软采样，可能会导致解码器收到不明确的边权重，轨迹预测反而变差。这个差异我自己的实测是：预测阶段用argmax的边恢复准确率更高，轨迹误差也略好。

3. 从零实现：弹簧质点系统上的NRI训练全流程

3.1 数据生成：制造一个真实连接已知的物理世界

要验证NRI能不能推理出结构，我们需要一个真实连接关系可以查、且动力学足够清晰的系统。弹簧质点系统是教科书级的测试床。代码思路如下：

• 固定N个质点在二维空间内运动。
• 每一对质点之间有概率p连接一根轻弹簧，弹簧自然长度为0，劲度系数为k。
• 每个样本随机初始化各质点的位置和速度。
• 用半隐式欧拉法做数值积分，生成一段足够长的轨迹。
• 把整条轨迹用滑窗切成“输入9帧 + 预测40帧”的训练样本。

这里有一个工程细节值得说明：物理仿真必须用向量化实现，否则在生成几万条样本的时候会慢到怀疑人生。对每一帧，计算所有质点两两之间的位移向量diff（形状[N, N, 2]），再乘上邻接矩阵得到每个质点收到的弹簧力：

import torch def simulate_trajectory(adj, num_steps=60, dt=0.1, k=0.2): N = adj.shape[0] pos = torch.rand(N, 2) * 2.0 - 1.0 vel = (torch.rand(N, 2) - 0.5) * 0.1 traj = [] for _ in range(num_steps): diff = pos[:, None, :] - pos[None, :, :] # [N, N, 2] force = adj[..., None] * diff * k # 线性弹簧，自然长度0 acc = force.sum(dim=0) # 每个节点受合力 vel = vel + acc * dt pos = pos + vel * dt traj.append(torch.cat([pos, vel], dim=-1)) return torch.stack(traj, dim=0) # [T, N, 4]

这个实现里用到的线性弹簧力模型是F = k·r，方向指向连接的另一端。它的物理含义是：两个质点离得越远，拉力越大。这样的简化让代码非常简洁，同时动力学足够非线性，能体现出NRI相对基线的优势。

3.2 模型搭建：Encoder和Decoder的PyTorch实现

数据准备好了，接下来是核心模型。我把编码器的实现写在这里，整体思路是“节点特征降维 → 边消息传递 → 节点更新 → 边类型logits”。注意代码里有两个mask的地方，一个是消息聚合时屏蔽自环，一个是输出时屏蔽对角线，这点非常重要，否则模型会出现“自己和自己有边”的幻觉。

import torch import torch.nn as nn import torch.nn.functional as F class MLP(nn.Module): def __init__(self, in_dim, hidden_dim, out_dim, n_layers=2): super().__init__() layers = [] for i in range(n_layers): in_d = in_dim if i == 0 else hidden_dim out_d = out_dim if i == n_layers - 1 else hidden_dim layers.append(nn.Linear(in_d, out_d)) if i < n_layers - 1: layers.append(nn.ReLU()) self.net = nn.Sequential(*layers) def forward(self, x): return self.net(x) class NRIEncoder(nn.Module): def __init__(self, feat_dim=4, hidden_dim=128, edge_types=4): super().__init__() self.node_embed = nn.Linear(feat_dim, hidden_dim) self.edge_mlp_1 = MLP(hidden_dim * 2, hidden_dim, hidden_dim) self.node_mlp_1 = MLP(hidden_dim * 2, hidden_dim, hidden_dim) self.edge_mlp_2 = MLP(hidden_dim * 2, hidden_dim, hidden_dim) self.logit_fc = nn.Linear(hidden_dim, edge_types) def forward(self, x): # x: [B, T, N, F] B, T, N, F = x.shape h = self.node_embed(x).mean(dim=1) # [B, N, H] # 构造边的特征，h_i和h_j分别广播 h_i = h.unsqueeze(2).expand(-1, -1, N, -1) h_j = h.unsqueeze(1).expand(-1, N, -1, -1) edge_feat = torch.cat([h_i, h_j], dim=-1) # [B, N, N, 2H] e = self.edge_mlp_1(edge_feat) # [B, N, N, H] # 聚合邻居消息，屏蔽自环 mask = torch.eye(N, dtype=torch.bool, device=x.device) agg = e.masked_fill(mask.unsqueeze(0).unsqueeze(-1), 0).sum(dim=2) h_new = self.node_mlp_1(torch.cat([h, agg], dim=-1)) # 第二轮边更新，输出logits h_i2 = h_new.unsqueeze(2).expand(-1, -1, N, -1) h_j2 = h_new.unsqueeze(1).expand(-1, N, -1, -1) edge_feat2 = torch.cat([h_i2, h_j2], dim=-1) e2 = self.edge_mlp_2(edge_feat2) logits = self.logit_fc(e2) # [B, N, N, K] # 对角线置为极小的logits，确保无自环 logits = logits.masked_fill(mask.unsqueeze(0).unsqueeze(-1), -1e9) return logits

编码器里值得注意的点是两次消息传递。第一轮把节点特征映射到边空间，聚合邻居信息后更新节点表示；第二轮再基于更新后的节点表示计算边logits。这样处理之后，logits不仅反映了两个节点的个体状态，还包含了它们各自邻居的信息，对局部运动模式的感知会更充分。

解码器按前面讲的逻辑实现。这里我用了一个小技巧：GRU的输入不只是邻居消息聚合成的结果，还会把当前状态（位置和速度拼接）也通过一个线性层嵌入后加进去，让循环单元能感知当前的物理状态。

class NRIDecoder(nn.Module): def __init__(self, feat_dim=4, hidden_dim=128, edge_types=4, pred_len=40, dt=0.1): super().__init__() self.edge_types = edge_types self.pred_len = pred_len self.dt = dt self.node_embed = nn.Linear(feat_dim, hidden_dim) self.state_embed = nn.Linear(feat_dim, hidden_dim) self.edge_mlps = nn.ModuleList([ MLP(hidden_dim * 2, hidden_dim, hidden_dim) for _ in range(edge_types) ]) self.gru = nn.GRUCell(hidden_dim, hidden_dim) self.out_fc = nn.Linear(hidden_dim, 2) def forward(self, x, z): # x: [B, T0, N, F], z: [B, N, N, K] B, T0, N, F = x.shape pos = x[:, -1, :, :2].clone() vel = x[:, -1, :, 2:].clone() h = self.node_embed(x[:, -1]).reshape(B * N, -1) preds = [x[:, -1]] mask = torch.eye(N, dtype=torch.bool, device=x.device) for _ in range(self.pred_len): hn = h.reshape(B, N, -1) # 构造 [B, N, N, 2H] h_i = hn.unsqueeze(2).expand(-1, -1, N, -1) h_j = hn.unsqueeze(1).expand(-1, N, -1, -1) edge_feat = torch.cat([h_i, h_j], dim=-1) messages = torch.zeros(B, N, N, hn.shape[-1], device=x.device) for k in range(self.edge_types): msg_k = self.edge_mlps[k](edge_feat) # [B, N, N, H] z_k = z[..., k].unsqueeze(-1) # [B, N, N, 1] messages += z_k * msg_k # 聚合邻居，屏蔽自环 agg = messages.masked_fill(mask.unsqueeze(0).unsqueeze(-1), 0) agg = agg.sum(dim=2) # [B, N, H] state = torch.cat([pos, vel], dim=-1) # [B, N, F] state_emb = self.state_embed(state).reshape(B * N, -1) h = self.gru(agg.reshape(B * N, -1) + state_emb, h) dv = self.out_fc(h).reshape(B, N, 2) vel = vel + dv * self.dt pos = pos + vel * self.dt preds.append(torch.cat([pos, vel], dim=-1)) return torch.stack(preds, dim=1) # [B, pred_len+1, N, F]

这个解码器在实现上做了一个简化：原论文在每一步会用GRU的输出去计算下一刻的输入特征并再次嵌入，我这里把当前的位置和速度直接作为状态嵌入喂给GRU，效果上差异很小，但代码读起来清爽很多。有一点必须提醒，vel = vel + dv * self.dt中的dv其实是加速度的预测值，严格说应该叫acceleration而非velocity increment。在代码注释里写清楚就行，不影响可读性。

3.3 训练循环：ELBO损失与端到端优化

训练数据生成完成后，把每个样本的输入设置为前9帧，预测目标设置为后续40帧。训练时用Gumbel-Softmax从编码器得到的logits里采样软图结构，喂给解码器，再计算MSE重建损失和KL散度。

def train_step(model_enc, model_dec, opt, x_in, y_target, tau=1.0, kl_weight=1.0): opt.zero_grad() logits = model_enc(x_in) # [B, N, N, K] # Gumbel-Softmax采样 z_soft = F.gumbel_softmax(logits, tau=tau, hard=False, dim=-1) pred = model_dec(x_in, z_soft) # [B, pred_len+1, N, F] B, T, N, F = pred.shape loss_rec = F.mse_loss(pred[:, 1:], y_target) # KL散度: q(z|x) 对均匀先验 log_p = F.log_softmax(logits, dim=-1) kld = -log_p.mean(dim=-1).sum(dim=(1, 2)).mean() / N loss = loss_rec + kl_weight * kld loss.backward() opt.step() return loss.item(), loss_rec.item(), kld.item()

训练过程中可以实时观察边恢复准确率。因为数据集生成时保存了真实的邻接矩阵，测试时直接把编码器输出的logits做argmax，和真实邻接矩阵比对。我在下面章节给出完整的评估逻辑和代码片段。

超参数方面，我用的是Adam优化器，初始学习率5e-4，batch size 32，训练60个epoch。这样一个实验在单张普通显卡上大概几分钟就能跑完，CPU上会慢一些但也能接受。节点数N取10，每对边的连接概率0.5，边类型K取2（有弹簧和没弹簧），轨迹用滑窗切出30000个训练样本。

3.4 评估：结构恢复准确率和轨迹预测误差

测试评估需要做两件事。第一是结构恢复：把编码器输出的logits在最后一个维度上做argmax，得到预测的边类型，和真实邻接矩阵计算精确率、召回率和F1。第二是轨迹预测：和几个基线（比如直接用LSTM预测每个节点的未来轨迹）对比MSE。

def evaluate(model_enc, model_dec, test_loader): model_enc.eval() model_dec.eval() edge_correct = 0 edge_total = 0 total_mse = 0.0 with torch.no_grad(): for batch in test_loader: x_in, y_target, adj_true = batch logits = model_enc(x_in) z_pred = logits.argmax(dim=-1) # [B, N, N] # 无向图：取上三角 z_pred_triu = z_pred[:, torch.triu_indices(z_pred.size(1), z_pred.size(2), offset=1)[0], torch.triu_indices(z_pred.size(1), z_pred.size(2), offset=1)[1]] adj_triu = adj_true[:, torch.triu_indices(adj_true.size(1), adj_true.size(2), offset=1)[0], torch.triu_indices(adj_true.size(1), adj_true.size(2), offset=1)[1]] edge_correct += (z_pred_triu == adj_triu).sum().item() edge_total += adj_triu.numel() z_hard = F.one_hot(logits.argmax(dim=-1), num_classes=logits.shape[-1]).float() pred = model_dec(x_in, z_hard) total_mse += F.mse_loss(pred[:, 1:], y_target).item() * x_in.size(0) acc = edge_correct / edge_total mse = total_mse / len(test_loader.dataset) return acc, mse

我在实验中发现一个有意思的现象：当弹簧连接概率p设置为0.5时，NRI在测试集上的边恢复准确率可以到92%以上，也就是说模型成功找出了超过九成的弹簧连接。这个结果不是偶然，关键在于解码器的监督信号足够强——预测未来的任务迫使编码器必须找到那个能合理解释运动的图结构，否则多步预测会迅速偏离真实轨迹。这正体现了端到端优化的力量：结构不是靠额外标注学出来的，而是靠“预测效果”这条隐含的监督线逼出来的。

4. 复现NRI时绕不开的调参坑与边界问题

4.1 Gumbel温度：退火太快会让结构推理失效

温度tau是NRI里最值得花时间调的超参数。很多复现笔记里都有一个误区：温度必须从大往小退火，否则采样的梯度无法有效传递。我的实测结论是，在NRI这个框架里，固定tau=1.0训练往往比花哨的退火策略更省心。

原因分析：Gumbel-Softmax的分布本身是非对称的，温度偏大时softmax输出比较平滑，编码器的梯度更稳定，但解码器接收到的是一个软图，节点消息会在多种边类型之间加权融合，这反而让解码器对结构错误更鲁棒。如果把温度降得太低，编码器输出接近硬one-hot，梯度方差急剧增大，训练早期阶段解码器还没学会利用结构信息，这时一个错误的离散边会带来极大的loss波动，训练很容易震荡甚至不收敛。

如果你想尝试退火，建议采用余弦退火而不是线性衰减，并且在早期至少保持1000步的“预热期”。我个人最后的选择是固定1.0，实测在弹簧系统和行人轨迹数据上都稳定。

4.2 KL权重和先验失衡：不要被先验拖住

在物理系统里，“两个物体之间没有连接”通常占多数。如果先验设成均匀分布，而数据里只有20%的边存在，KL散度会施加一个很强的“别预测有边”的压力，模型会倾向把所有边都预测为“无连接”。这种情况下，KL权重需要降低，比如从1.0下调到0.1，或者干脆把先验改为类条件均匀分布，根据训练集里每类边的大致比例来设置先验频率。

另外一个更巧妙的做法是：对loss里的KL项按边类型做加权，把稀疏类别的KL惩罚减小。这对多类型边（比如有弹簧、排斥力、摩擦）尤其重要，因为稀少的交互类型如果受到过强先验压制，几乎不可能被模型恢复出来。

4.3 长轨迹预测的误差累积问题

NRI的解码器是自回归结构，训练时每一帧的输入都来自真实轨迹（teacher forcing），推理时上一帧的输出会作为下一帧的输入，所以误差会随着预测步数增长而累积。在40步预测以内，这个问题还不太明显，一旦预测长度超过100步，你会发现轨迹偏差呈指数级恶化。

缓解办法是引入一种类似课程学习的策略：训练时随机选择展开长度，而不是每次都直接展开到40步，让解码器逐步学会在自身预测上继续预测。具体来说，可以用一个范围在[1, 40]的随机整数，每次训练只在这个长度内做自回归展开，并把展开过程中解码器自己的预测拼接回输入作为下一步的初始状态。这个改动看起来小，但对长轨迹预测的提升非常明显。

4.4 图规模扩展性和置换不变性的隐含假设

NRI的编码器在构造边特征时，需要显式构造[B, N, N, ...]的边张量，时间复杂度是O(N²)，空间复杂度同理。当N达到一百甚至上千时，这种方法会直接爆显存。目前业界处理这个问题通常是分块计算边特征，或者用图采样技术只对部分邻居做消息传播，牺牲一部分精度换取可扩展性。

更要留意的是一个隐含假设：NRI假设系统是同构的，也就是所有节点的“身份”是等价的，交换任意两个节点，模型的输出概率不变。这在很多场景下不成立。比如行人轨迹预测中，不同人有不同的运动意图和个性；交通场景里车辆和行人本身就属于不同类型。如果你的系统里有明显的节点类型差异，最直接的办法是在输入特征里加一个one-hot的类型编码，或者在消息函数里把发送者和接收者的类型也作为输入的一部分，让模型学到类型相关的交互模式。

5. 从NRI出发能走多远：局限、扩展与替代思路

5.1 推理出的结构是“相关”而非“因果”

NRI找到的边类型本质上是在当前观测数据下、对运动变化最有解释力的结构关系，但并不自动等价于物理因果。举一个例子：如果两个质点同时被一个隐藏在暗处的力场驱动，它们的位置变化会高度相关，NRI很可能在它们之间推理出一条“弹簧连接”，而实际上并不存在直接的物理连接。这是一个典型的混淆因子问题。

所以在把NRI推理出的结构用于因果分析时，需要额外的干预实验或领域知识做交叉验证。在纯预测任务上这个短板不影响使用，但如果你的目标是“发现系统真正的连接关系”，务必在实验设计上引入对照。

5.2 动态交互：当关系本身随时间变化

NRI假设z_ij在整个观测序列和预测序列中保持不变。这个假设在多体物理系统里是合理的（弹簧连接不会突然消失），但在很多真实场景里站不住脚。比如两辆车在高速上并排行驶一段时间后分开，社交场景里两个人的交互有开始、有结束。

针对这类问题，后续工作提出了动态NRI（dNRI）：把潜变量从静态的z变成随时间演化的序列z_t，解码器在每个时间步都重新采样边结构。代价是优化复杂度上升，因为潜变量序列的推理需要用到类似变分序列推断的方法。如果你的数据里交互关系确实在时间上变化，建议优先考虑这类扩展。

5.3 和大模型、神经算子类方法的对比

近两三年，直接用Transformer做轨迹预测和用图神经网络做轨迹预测的赛道发生了交叉。Transformer的注意力机制天然能处理变长序列，但它是全局密集交互，计算复杂度和节点数量平方相关，而且注意力权重并不天然等于物理连接。NRI最大的优势是紧致的离散结构先验——它把问题从“对所有节点两两建模”压缩到“先找稀疏结构，再按结构做消息传递”，这带来了更高的数据效率和更好的泛化能力。

在样本量极少的情况下（比如只有几百条轨迹），Transformer几乎无法训练出有意义的结构，而NRI因为把结构作为一种显式归纳偏置注入模型，即使在小数据上也能恢复出大致正确的图结构。我的建议是：如果系统本身存在清晰的交互结构，优先使用NRI这类方法；如果交互模式高度复杂且难以用离散边类型表达，再考虑Transformer或混合架构。

5.4 几个值得继续探索的方向

从我复现和二次开发的经验来看，有几条路性价比很高。

第一条是把NRI作为因果发现工具，在多体动力学之外的应用上做迁移。比如把分子动力学模拟中的原子坐标输入NRI，看它能否恢复出化学键结构，这个方向在计算化学里已经有论文做过并显示出不错的结果。

第二条是引入物理先验，比如让解码器的消息函数尊重牛顿第三定律，也就是把消息函数约束成反对称的，或者直接使用更贴近真实物理的势能函数来参数化消息。这样不仅能提高预测精度，还能让推理出的边类型更有物理可解释性。

第三条是扩大节点规模。目前社区里有不少工作用图分区或图采样的方式把NRI扩展到几百节点，虽然不是官方实现，但代码也不复杂，值得研究。如果顺着这条路径发展，NRI的适用范围会从“小规模多体系统”拓展到“大规模社交网络或城市交通网络”级别的问题。

我在实际跑NRI的过程里最大的体会是，这个模型的思路并不过时，它的精髓在于“先推理结构再做预测”这个归纳偏置，在数据量有限的物理问题中，这个偏置远比模型容量重要。如果你正准备拿图神经网络处理轨迹预测或者结构发现类任务，NRI是绕不过去的起点。顺着它的思路，你可以很自然地把注意力机制、动态潜变量、物理约束等现代技术嫁接到这个框架上，做出更符合自己场景的解决方案。