NeuCLIP: Efficient Large-Scale CLIP Training with Neural Normalizer Optimization

TL;DR

NeuCLIP introduces a novel neural normalizer optimization framework for large-scale CLIP training, significantly improving the estimation of normalization terms and enhancing model performance on massive datasets.

Contribution

It reformulates the contrastive loss into a minimization problem with an auxiliary variable and transforms it into a neural network prediction task, enabling more accurate normalization in large-scale CLIP training.

Findings

Outperforms previous normalizer estimation methods on datasets from millions to billions of samples.

Achieves more accurate normalization leading to improved CLIP model performance.

Demonstrates scalability and efficiency in large-scale contrastive learning.

Abstract

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.