TensorFlow自定义消息传递GNN拼接边属性时梯度爆炸问题求解

为什么取值范围正常的e会触发梯度爆炸 #

虽然e本身取值仅在0.5~4区间，不存在异常值，但核心原因是特征尺度不匹配引发的权重偏置更新：

• 你用来拼接的neighbors_mean是节点属性的均值，通常节点属性在输入GNN前会做归一化，分布在0附近的小区间（比如[-1,1]、[0,1]），而e的数值区间和节点属性不一致，拼接后e对应的特征维度在单层前馈网络Net的计算中会主导输出结果，反向传播时对应e维度的权重梯度会远大于其他维度的权重梯度。
• 多轮训练后，Net中对应e维度的权重数值会被更新得越来越大，再叠加GNN消息传递层的梯度反向传播累积效应（梯度跨多层消息传递层连乘），最终就会出现梯度爆炸。你之前用sigmoid压缩e的操作本质是把e的尺度对齐到节点特征的区间，所以能避免爆炸，但非线性变换丢失了距离的原始线性关系，才导致MAE上升。

可尝试的解决方案（不会损失e的有效信息） #

• 对e做线性归一化对齐特征尺度：直接用线性变换把e缩放到和neighbors_mean完全一致的数值区间，完全保留距离的相对关系，不会损失信息。比如neighbors_mean归一化到[-1,1]，就用如下代码处理e：

# e原始范围0.5~4，线性缩放到[-1,1]
e_norm = 2 * (e - 0.5) / (4 - 0.5) - 1
to_concat = [neighbors_mean, e_norm]

• 拼接前对e做线性投影：给e单独加一个不带激活的线性层，把e投影到和节点特征相同的隐空间，自动对齐尺度：

# 层初始化仅需执行一次，不要每次调用都新建层
self.e_proj = tf.keras.layers.Dense(units=8, activation=None) # units可根据节点特征维度调整
e_proj = self.e_proj(e)
to_concat = [neighbors_mean, e_proj]

• 拼接后增加层归一化：在z输入Net之前加一层层归一化，强制所有特征维度的分布统一为0均值1方差，不会破坏e的相对大小信息：

z = K.concatenate(to_concat, axis=-1)
z = tf.keras.layers.LayerNormalization()(z)
output = self.Net(z)

• 限制Net输出的数值范围：如果你的Net用了ReLU等无上界的激活函数，可以给输出加一个宽松的数值裁剪，避免前向计算溢出：

output = tf.clip_by_value(self.Net(z), -15, 15)

内容的提问来源于stack exchange，提问作者Achintha Ihalage