💻 computer science

Memory-Efficient Contrastive Learning via Budgeted Hard Negative Selection

本文介绍了一种内存高效的对比学习框架，该框架通过流式计算并动态选择固定预算的难负样本，消除了稠密相似度矩阵带来的二次方内存瓶颈，从而在保持优化有效性的同时，使受限硬件能够支持显著更大的批次大小。

原作者： Qinwu Xu

发布于 2026-09-07

📖 1 分钟阅读☕ 轻松阅读

原作者： Qinwu Xu

原始论文采用 CC BY 4.0 许可（https://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

在人工智能领域，计算机正日益通过将图像进行相互比较来学习如何观察和理解世界。想象一名学生正在学习什么是“狗”的样子。他不是被告知一个定义，而是被展示了数千张图片。为了有效地学习，这名学生不仅要识别出两张狗的照片是相似的，还要理解一张狗的照片与一张汽车或一棵树的照片之间有多么不同。这个过程被称为对比学习（contrastive learning），它是许多现代视觉系统的引擎。它的工作原理是将相似的事物在数学空间中拉近，同时将不同的事物推开。系统一次看到的示例越多，且能越清晰地分辨它们，它就会变得越聪明。然而，计算机在任何单一时刻能存储在内存中的信息量存在物理极限。随着研究人员试图向这些系统喂入越来越大的图像批次以改进学习，计算机的内存经常会写满并崩溃，就像一个背包在你试图往里面塞进最后一本重书时崩裂了一样。

德克萨斯大学奥斯汀分校的一位研究人员开发了一种运行这些学习系统的新方法，避免了这种内存爆炸。他们的研究方法详述于 2026 年 9 月发表的一项研究中，该方法改变了计算机处理用于教学系统所需的庞大比较列表的方式。传统上，为了比较一组图像，计算机会创建一个巨大的网格，同时计算每一张图像与所有其他图像之间的相似性。如果一组图像包含四千张，这个网格将需要数百万次计算，并且需要大量的内存来存储这些数字。研究人员发现，虽然计算机需要知道图像之间的确切关系来进行学习，但它并不需要将整个网格同时保存在内存中。相反，他们设计了一种方法，将这些比较过程分解为细小、可控的数据块，通过系统进行数据流式传输，而不是囤积数据。

这种新方法的核心是一种被称为“预算化硬负样本选择”（budgeted hard negative selection）的技术。在学习过程中，并非所有的差异都同样重要。有些图像与目标对象差异巨大，以至于计算机无法从中学习到任何新知识；这些被称为“易负样本”（easy negatives）。而另一些图像虽然非常相似但又不完全相同，这些才是真正驱动学习的“硬负样本”（hard negatives）。新系统将注意力集中在寻找这些困难且具有信息量的示例上，同时忽略那些容易的示例。它是通过对图像进行小块处理来实现这一点的。当它处理每个数据块时，它会保留一份目前为止发现的最难示例的运行列表。如果一个新的数据块到达，且其中的图像都不比现有列表中的图像更难，系统就会直接跳过排序和存储的工作。这就像一位图书管理员，在将新书与热门书单进行对比时，只有当发现一本新书比当前列表中最不热门的书更受欢迎时，才会停下来更新列表；否则，这本书只是被扫一眼后就被搁置在一旁。

通过使用这种流式传输方法，研究人员能够大幅减少训练这些模型所需的内存。在测试中，他们使用了一块拥有 80 GB 显存的强大显卡。标准方法在批次大小达到 4,096 张图像时就会耗尽内存。然而，新方法在相同的硬件上成功训练了 8,192 张图像的批次。比较数据的内存占用从“二次方增长”（即图像数量翻倍，内存需求变为原来的四倍）降到了“线性增长”（即图像数量翻倍，内存需求也仅随之翻倍）。这使得系统能够处理两倍数量的示例而不发生崩溃。此外，随着训练的推进，系统变得更加高效。到第十轮训练时，近 90% 的潜在比较都被跳过了，因为系统已经找到了更好的示例，从而节省了大量的处理时间。

为了使系统更加精简，研究人员将这种流式传输法与另外两种效率工具相结合。一种工具使用队列来存储来自前几轮训练的示例，使系统能够在不需要将所有图像都保存在活动内存中的情况下，从更广泛的图像中学习。另一种被称为“低秩自适应”（low-rank adaptation）的工具，改变了计算机更新其内部知识的方式。它不再重写模型的整个庞大“大脑”，而只是调整一小组专门的参数。这种组合使得研究人员能够在单块显卡上训练复杂的视觉模型，在测试的最大批次下，内存占用低至 6.1 GB。该研究证实，这种方法并不会牺牲学习质量；通过这种方式训练的模型仍然能产生高质量的表示，并在标准的图像识别任务中表现出色。

研究人员强调，他们的方法并没有跳过比较图像的实际数学运算；它仍然计算每一对图像之间的确切相似性以确保准确性。创新之处完全在于如何存储和管理这些数据。通过拒绝将庞大的比较网格实体化，转而以稳定的流式传输方式处理数据，他们消除了训练大规模视觉系统的一个主要瓶颈。这项工作为在现有硬件上训练更智能、更强大的模型提供了实用的基础，证明了效率可以通过更智能地组织工作流程来实现，而不是通过在学习过程中偷工减料。结果表明，人工智能训练的限制往往取决于我们如何管理资源，而不仅仅取决于机器的原始动力。

技术摘要：通过预算化硬负样本选择实现内存高效的对比学习

问题陈述

对比学习框架（如 CLIP 和 SimCLR）依赖于大规模负样本集来学习共享嵌入空间。然而，传统的实现方式会为大小为 $N$ 的批次生成稠密的成对相似度矩阵（ $N \times N$ ）。这导致内存占用呈二次方增长（ $O(N^2)$ ），从而产生显著的系统级瓶颈。虽然评估所有成对点积对于精确发现硬负样本（hard-negatives）非常有用，但保留完整的矩阵是不必要的，因为训练目标（InfoNCE）最终仅利用一小部分具有信息量的“硬”负样本。此外，负样本的 Softmax 加权具有高度非均匀性：容易识别的负样本对优化贡献微乎其微，但它们在完整矩阵中的存在却消耗了大量的 GPU 显存。这一限制制约了在受限加速器（如单 GPU）上扩展批次大小的能力。

方法论

作者提出了一种内存高效的稀疏对比学习框架，该框架将相似度评估与相似度物化（materialization）解耦。其核心方法涉及分块流式处理（blockwise streaming）和预算化硬负样本选择，并辅以参数高效的适配技术。

1. 分块流式处理与稀疏挖掘

该方法并非计算完整的 $N \times N$ 相似度矩阵 $S = ZZ^\top$ ，而是以大小为 $b \ll N$ 的列块形式计算相似度。

流式计算： 相似度按顺序计算为 $S_{local} = Z \cdot Z_{block}^\top$ 。通过利用稠密矩阵-矩阵乘法（GEMM）而非元素级操作，这种方式保持了 GPU 的计算效率。
边界掩码（Boundary Masking）： 在每个块内，无效对（自对及正样本对）被掩码为 $-\infty$ 。
在线 Top-k 剪枝： 系统为每个锚点（anchor）维护一个运行中的前 $k$ $k$ 个硬负样本缓冲区（ $V_{run}$ $V_{r u n}$ ）。在针对特定行（锚点）在块内执行局部 top- $k$ $k$ 排序之前，会应用行级最大值测试：
- 如果当前块中锚点 $i$ 的最大相似度 $m_i$ 小于或等于当前第 $k$ 个硬负样本（ $V_{run}[i, -1]$ ），则该块无法改进候选集。
- 因此，该行的局部 top- $k$ 更新和合并操作将被跳过。
- 随着训练的进行以及硬负样本阈值的提高，这一剪枝规则消除了不必要的排序和内存操作。

2. 辅助效率机制

基于队列的记忆库（Queue-Based Memory Bank）： 为了在当前批次之外扩展负样本的多样性，一个脱离梯度的队列存储了来自前序步骤的嵌入。这些嵌入以非微分方式检索，用以扩充负样本集，而不会增加梯度存储成本。
LoRA 适配： 该框架集成了低秩适配（LoRA）以减少可训练状态的内存占用。通过冻结骨干网络并仅训练低秩矩阵（ $A$ 和 $B$ ），与全量微调相比，优化器状态的内存显著降低。

3. 复杂度与权衡

内存： 峰值相似度缓冲区存储从 $O(N^2)$ 降至 $O(Nb)$ 。持久化的硬负样本状态仅需 $O(Nk)$ 。
计算： 该方法扫描所有的列块，以确保实现精确的批次内 top- $k$ 负样本发现。因此，成对点积的最坏情况算术复杂度仍为 $O(N^2D)$ 。效率的提升严格体现在内存减少以及排序/合并操作的剪枝上，而非亚二次方算术复杂度的降低。

核心贡献

公式化： 作者将对比硬负样本选择建模为一个固定预算的流式问题，避免了稠密 $N \times N$ 相似度矩阵的物化。
剪枝准则： 引入了一种在线行级剪枝准则，当一个相似度块无法改进锚点的运行候选集时，可以安全地跳过局部 top- $k$ 更新。
集成流水线： 该方法结合了流式批次内挖掘、可选的脱离梯度记忆库以及基于 LoRA 的适配，为内存受限的场景创建了一个实用的流水线。
经验特征分析： 论文提供了关于内存-选择权衡的经验数据，包括批次大小扩展极限以及剪枝有效性在训练过程中的演变。

实验结果

实验在单张 NVIDIA A100 (80GB) GPU 上进行，使用 ViT-B/32 骨干网络，维度 $D=512$ 。

内存可扩展性：
- 稠密基准线在批次大小为 4096 时触发显存溢出（OOM）错误。
- 流式配置可扩展至批次大小 8192。
- 在 $N=2048$ 时，峰值内存分配从 48.2 GB（稠密）降至 11.4 GB（流式）。
- 结合 LoRA 和梯度检查点（gradient checkpointing），该配置在 $N=8192$ 时仅使用 6.1 GB 内存进行训练。
剪枝效率：
- 行级剪枝在训练轮数（epoch）增加的过程中变得越来越有效。
- 到第 10 轮时，89.1% 的行在流式块内跳过了局部 top- $k$ 更新。
- 这使得与未剪枝过程相比，top- $k$ 选择的工作量（排序和合并）减少了 3.62 倍，尽管成对 GEMM 算术运算保持不变。
表示质量：
- 通过在 CIFAR-100 上进行的标准线性评估协议，论文验证了稀疏训练过程保留了表示质量，尽管提供的文本中未详细说明具体的准确率数值。

意义与主张

本文为受限加速器内存约束下的大规模批次对比适配奠定了实践基础。其主要意义在于，能够在不牺牲批次内硬负样本发现精确性的前提下，实现目前在单台高端 GPU 上通过标准稠密实现无法达到的批次大小（例如 8192）。

作者明确指出，目前的实现优先考虑精确的硬负样本发现，并不声称具备亚二次方算术复杂度。效率的提升是通过内存减少（ $O(N^2) \to O(Nb)$ ）以及通过动态剪枝消除冗余排序操作来实现的。未来的工作确定为研究如何通过近似候选生成机制来完全避免评估无信息相似度块。