Amazon ML Coding：PyTorch InfoNCE，正样本下标与梯度怎样核对？

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• queries/keys同形[N,D]、同device/dtype，i与i配对正样本，其余key是本batch负样本；这是paired双塔版本，不是SimCLR的2N自相似矩阵。
• 温度τ有限且>0，默认L2normalize；symmetric=True取双向平均。输入有限，零向量由eps合同处理，但缺少语义方向应监控。
• 训练float16/bfloat16先升float32算normalize/logits/CE，float64保留以做数值验证；仅CPU实际验证，不声称已测GPU/DDP。

解题路线

1. 归一化q/k后score=qkᵀ/τ，shape[N,N]。第i行target=i，loss是−score_ii+logsumexp(score_i*)；对角是正确候选必须保留，不能套自对比mask把正样本删掉。
2. 直接给cross_entropy原始logits与long labels，不先softmax再CE；logsumexp稳定性避免大logits的exp溢出。labels放同device，最后loss标量mean。
3. 需要对称时同一score转置再CE，query→key与key→query各算平均；单向目标也合法但语义不同。N=1只有正样本，loss0无负样本对比信号，不能凭loss低说表示已学好。
4. 若不同样本其实同语义/同实体，off-diagonal会成假负样本；多positive需要不同target集合/目标，不能仅把每行target=i继续声称处理了多正样本。
5. 跨卡扩大候选需全局positive offset、all-gather是否带梯度、是否每卡算loss/重复归约等合同；普通不带grad gather可能阻断远端key梯度，单机代码不自动扩为正确DDP。
6. 独立验证先逐向量标量normalize/dot/logsumexp对照前向，再有限差分标量参考loss核两输入梯度；用不同temperature/N/D、零向量和低精度升型核边界。

手动走一遍

N2时q0对应k0，第一行候选k0/k1，target0；如果调换keys但不调positive索引，训练会把真正正样本当负。测试1080前向矩阵和1440梯度分量，加零向量、3种dtype与6非法输入通过；运行环境为临时隔离PyTorch CPU环境，版本记入审阅台账。

Python 参考实现

下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。

代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。

import math
import torch
import torch.nn.functional as F


def info_nce(queries, keys, temperature=0.1, symmetric=True):
    """Paired [N,D] embeddings; diagonal positives, other keys in-batch negatives."""
    if queries.ndim != 2 or queries.shape != keys.shape or min(queries.shape) == 0:
        raise ValueError("matching nonempty [N,D] embeddings required")
    if queries.device != keys.device or queries.dtype != keys.dtype:
        raise ValueError("matching device and dtype required")
    if not queries.is_floating_point() or not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("floating embeddings and finite positive temperature required")
    if not torch.isfinite(queries).all() or not torch.isfinite(keys).all():
        raise ValueError("finite embeddings required")
    # Preserve float64 for numerical checking, promote lower precision for stability.
    dtype = torch.float64 if queries.dtype == torch.float64 else torch.float32
    q = F.normalize(queries.to(dtype), dim=-1, eps=1e-12)
    k = F.normalize(keys.to(dtype), dim=-1, eps=1e-12)
    logits = (q @ k.T) / temperature
    labels = torch.arange(len(q), device=q.device)
    forward = F.cross_entropy(logits, labels)
    if not symmetric:
        return forward
    return (forward + F.cross_entropy(logits.T, labels)) / 2

复杂度

score O(N²D)时间、O(N²)空间，normalize O(ND)。大batch负样本多但矩阵与encoder激活更大，可分块CE但需保持全候选logsumexp和正确梯度，不按块各softmax当等价。

易错点与追问

容易失分的边界

• paired对角mask掉。
• softmax后再CE。
• labels在CPU而logits在GPU。
• 双向只写同一行CE两次。
• N1低loss当学习成功。
• 跨卡gather丢梯度。

继续深挖

• 多个positive如何定义分子？
• 分块logsumexp如何保数值稳定？
• false negatives怎样识别与缓解？
• DDP每卡positive下标怎样算？