TensorFlow中如何将三重嵌套循环转为并行张量运算优化损失函数？

问题背景 #

需要实现如下损失函数：
$$l_r = -\frac{2}{V(V-1)} \sum_{j=1}^C \sum_{k=1}^V \sum_{\substack{l=1 \ l \neq k}}^V \frac{\langle M_{j,l}, M_{j,k} \rangle}{|M_{j,l}|2 |M{j,k}|_2}$$
其中M是形状为576×2×2048的张量，C=576，V=2。原实现采用三重嵌套for循环，导致每个epoch运行效率极低，需要改用张量运算优化。

原低效代码：

M = self.kernel

norm_M = tf.norm(M, ord=2, axis=2)
norm_X = tf.norm(X, ord=2, axis=1)

# Compute reunion loss
sum = 0.0
for j in tf.range(C):
  for k in tf.range(V):
    for l in tf.range(V):
      if k == l:
        continue
      A = tf.tensordot(M[j][l], M[j][k], 1)
      B = norm_M[j][l] * norm_M[j][k]
      sum += A / B
l_r = -2/(self.V*(self.V-1)) * sum

张量运算优化方案 #

利用TensorFlow的批量矩阵运算和广播机制，完全消除Python循环，实现GPU加速的高效计算：

M = self.kernel  # 形状：[C, V, D]，C=576，V=2，D=2048

# 1. 对所有向量做L2归一化，归一化后向量点积直接等于余弦相似度
norm_M = tf.norm(M, ord=2, axis=2, keepdims=True)  # 形状：[C, V, 1]，保留维度以便广播
normalized_M = M / norm_M  # 形状：[C, V, D]，广播完成逐向量归一化

# 2. 批量计算每个j对应的V×V余弦相似度矩阵
# 通过矩阵乘法实现批量点积：[C,V,D] × [C,D,V] → [C,V,V]
cos_sim_matrix = tf.matmul(normalized_M, normalized_M, transpose_b=True)

# 3. 过滤对角线元素（k=l的情况）并求和
mask = 1 - tf.eye(V, dtype=tf.float32)  # 构造掩码，对角线为0，其余为1，形状：[V,V]
masked_cos_sim = cos_sim_matrix * mask  # 保留非对角线的余弦相似度值
total_sum = tf.reduce_sum(masked_cos_sim)  # 对所有非对角线元素求和

# 4. 计算最终损失
l_r = -2 / (self.V * (self.V - 1)) * total_sum

优化说明 #

• 完全移除三重循环，改用TensorFlow原生的批量张量运算，充分利用GPU并行计算能力
• 归一化后直接通过矩阵乘法批量计算所有余弦相似度，避免了逐元素循环的Python overhead
• 掩码操作高效过滤掉k=l的无效项，替代原代码中的条件判断

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