题目内容

实现 InfoNCE（对比学习）损失函数的计算。

给定一个 batch 的 $n$ 个样本，每个样本有两个向量表示（正对 pair）：anchor $a_i$ 和 positive $p_i$，维度为 $d$。计算 InfoNCE 损失：

对每个 anchor $a_i$，其正样本为 $p_i$，其余 $n-1$ 个 $p_j$（$j \ne i$）为负样本。相似度为余弦相似度除以温度参数 $\tau$：

输入描述

第一行包含两个整数 $n, d$（$2 \le n \le 64$，$1 \le d \le 32$）和一个浮点数 $\tau$（$0.01 \le \tau \le 1.0$）。

接下来 $n$ 行，每行 $d$ 个浮点数，表示 anchor 矩阵 $A$。

接下来 $n$ 行，每行 $d$ 个浮点数，表示 positive 矩阵 $P$。

输出描述

输出 InfoNCE 损失值，保留 $4$ 位小数。

样例 $1$

输入

3 2 0.5
1.0 0.0
0.0 1.0
1.0 1.0
1.0 0.0
0.0 1.0
1.0 1.0

输出

0.6000

说明

$3$ 个样本，温度 $\tau = 0.5$。每个 anchor 与自己的 positive 是正对，余弦相似度为 $1.0$（自身匹配），缩放后为 $2.0$。与其他 positive 的相似度取决于向量夹角。