如何基于PyTorch实现可微分的MaxSim算子并生成得分矩阵？

问题分析 #

给定两个张量批次：

• T：形状为(4, 3, 2)，包含4个3行2列的矩阵
• L：形状为(4, 3, 2)，结构与T一致

需要生成得分矩阵S（形状(4, 4)），其中S[i,j]是T中第i个矩阵与L中第j个矩阵的MaxSim得分，且整个计算过程完全可微分。

MaxSim核心逻辑：对两个矩阵的行向量两两计算相似度，取所有相似度中的最大值作为最终得分。

代码实现（PyTorch） #

import torch

def max_sim_score(T, L, sim_func=None):
    # 默认使用余弦相似度作为行向量相似度计算函数
    if sim_func is None:
        def sim_func(a, b):
            # a: (N, D), b: (M, D) -> 返回(N, M)的相似度矩阵
            a_norm = a / a.norm(dim=1, keepdim=True)
            b_norm = b / b.norm(dim=1, keepdim=True)
            return torch.matmul(a_norm, b_norm.T)
    
    # T: (B1, N, D), L: (B2, M, D)
    B1, N, D = T.shape
    B2, M, D = L.shape
    
    # 扩展维度，实现批次间两两配对计算
    T_expanded = T.unsqueeze(1)  # 形状: (B1, 1, N, D)
    L_expanded = L.unsqueeze(0)  # 形状: (1, B2, M, D)
    
    # 计算所有行对的相似度，重塑维度简化批量计算
    sim_matrix = sim_func(T_expanded.view(B1*B2, N, D), L_expanded.view(B1*B2, M, D))
    sim_matrix = sim_matrix.view(B1, B2, N, M)
    
    # 取每个批次配对的最大相似度
    max_sim = sim_matrix.max(dim=2)[0].max(dim=2)[0]
    
    return max_sim

# 示例测试
T = torch.randn(4, 3, 2, requires_grad=True)
L = torch.randn(4, 3, 2, requires_grad=True)

S = max_sim_score(T, L)
print("得分矩阵S形状:", S.shape)  # 输出: torch.Size([4, 4])

# 验证可微分性
loss = S.sum()
loss.backward()
print("T的梯度是否存在:", T.grad is not None)  # 输出: True
print("L的梯度是否存在:", L.grad is not None)  # 输出: True

关键细节说明 #

• 维度扩展：通过unsqueeze实现T和L的批次两两配对，避免循环计算，提升效率。
• 可微分保证：所有操作均使用PyTorch原生张量函数，包括norm、matmul、max，均支持自动求导。
• 灵活性：支持传入自定义的行相似度函数（如点积、归一化欧氏距离等），只需保证函数本身可微分即可。

内容的提问来源于stack exchange，提问作者Celso França