用神经网络优化对比学习归一化项,让大模型训练更准更快
NeuCLIP: Efficient Large-Scale CLIP Training with Neural Normalizer Optimization
- 将每个样本的归一化项建模为可学习的神经网络,避免传统近似误差
- 在百万到十亿级数据上训练,性能超越现有方法,尤其适合小批量场景
- 适合做大规模图文预训练的团队,特别是资源受限但追求精度的研究者
准确估计对比损失中的归一化项(即分区函数)是训练对比语言-图像预训练(CLIP)模型的核心挑战。传统方法依赖大批次近似,计算开销高。先前工作引入逐样本归一化器,在每轮迭代中分块更新以跟踪编码器变化,但其优化误差随数据集大小与批次大小之比增长,限制了在大数据集或小批次下的效果。为此,我们提出NeuCLIP,一种新颖且简洁的优化框架:(i) 基于凸分析,将每个样本的对比损失重构成含辅助变量(表示对数归一化项)的最小化问题;(ii) 通过变分分析,将对n个辅助变量(n为数据集规模)的最小化转化为对紧凑神经网络的最小化,该网络预测所有样本的对数归一化项。设计交替优化算法,联合训练CLIP模型与辅助网络。通过定制架构与加速技术,NeuCLIP实现更精确的归一化估计,显著提升性能。在涵盖百万至十亿样本的大规模CLIP训练中,实验表明其优于以往方法。代码已开源:https://github.com/Optimization-AI/NeuCLIP。
原文摘要 · Abstract (English)
Accurately estimating the normalization term (also known as the partition function) in the contrastive loss is a central challenge for training Contrastive Language-Image Pre-training (CLIP) models. Conventional methods rely on large batches for approximation, demanding substantial computational resources. To mitigate this issue, prior works introduced per-sample normalizer estimators, which are updated at each epoch in a blockwise coordinate manner to keep track of updated encoders. However, this scheme incurs optimization error that scales with the ratio of dataset size to batch size, limiting effectiveness for large datasets or small batches. To overcome this limitation, we propose NeuCLIP, a novel and elegant optimization framework based on two key ideas: (i) $\textbf{reformulating}$ the contrastive loss for each sample $\textbf{via convex analysis}$ into a minimization problem with an auxiliary variable representing its log-normalizer; and (ii) $\textbf{transforming}$ the resulting minimization over $n$ auxiliary variables (where $n$ is the dataset size) via $\textbf{variational analysis}$ into the minimization over a compact neural network that predicts the log-normalizers. We design an alternating optimization algorithm that jointly trains the CLIP model and the auxiliary network. By employing a tailored architecture and acceleration techniques for the auxiliary network, NeuCLIP achieves more accurate normalizer estimation, leading to improved performance compared with previous methods. Extensive experiments on large-scale CLIP training, spanning datasets from millions to billions of samples, demonstrate that NeuCLIP outperforms previous methods. Code is available at https://github.com/Optimization-AI/NeuCLIP.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。