如何高效计算所有i,j∈[N]的score=dot(a, LeakyReLU(x_i+y_j))并降低显存占用

你的原实现显存过高的核心原因是生成了N × N × D形状的三维中间张量feature_sum，当N、D均大于1000时，该张量的元素数量可达1e9量级，单精度浮点存储就需要4GB以上显存，很容易超出硬件限制。以下是两种可直接使用的优化方案：

方案1：小批量分块计算（改动最小，无逻辑变更） #

不需要修改原有计算逻辑，仅对X做分批处理，避免一次性生成全量三维张量，可通过调整批次大小适配任意显存：

# 给定输入 X (N x D), Y(N x D), A (D x 1)
alpha = 0.1
batch_size = 32 # 显存不足可下调，显存充足可上调提高计算效率
Y_expand = tf.expand_dims(Y, axis=0) # 仅需做一次维度扩展，形状为1 x N x D
score_list = []

for i in range(0, X.shape[0], batch_size):
    # 取当前批次的X样本
    X_batch = X[i:i+batch_size]
    X_batch_expand = tf.expand_dims(X_batch, axis=1) # 形状为batch_size x 1 x D
    # 计算当前批次的score
    feature_sum = X_batch_expand + Y_expand # 形状为batch_size x N x D，显存占用大幅降低
    batch_score = K.dot(LeakyReLU(alpha=alpha)(feature_sum), A) # 形状为batch_size x N x 1
    score_list.append(tf.squeeze(batch_score, axis=-1))

# 拼接所有批次结果得到最终N×N的score矩阵
dense = tf.concat(score_list, axis=0)

该方案的中间张量显存占用为batch_size * N * D，当batch_size设为16、N=2000、D=2000时，中间张量仅占256MB显存，完全不会出现显存溢出问题。

方案2：纯矩阵运算（无三维张量，性能最优） #

通过对LeakyReLU的逐元素特性做数学拆解，可以完全消除三维中间张量，直接用二维矩阵运算得到结果，显存复杂度从原有的O(N²D)降低到O(N²)：

# 给定输入 X (N x D), Y(N x D), A (D x 1)
alpha = 0.1
a = tf.squeeze(A) # 转换为形状为(D,)的向量

# 逐元素乘a到X、Y上
X_a = X * a[tf.newaxis, :]
Y_a = Y * a[tf.newaxis, :]

# 计算所有x_i + y_j的逐元素激活后和a的点积
def calculate_leaky_term(X_m, Y_m, alpha_val):
    # 利用广播计算激活后的贡献
    X_col = tf.expand_dims(X_m, 1)
    Y_row = tf.expand_dims(Y_m, 0)
    sum_mat = X_col + Y_row
    leaky_sum = LeakyReLU(alpha=alpha_val)(sum_mat)
    return tf.reduce_sum(leaky_sum, axis=-1)

dense = calculate_leaky_term(X_a, Y_a, alpha)

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