Research map / Architectures and efficiency / RNNs and normalization

Breaking the Memory Barrier of Contrastive Loss via Tile-Based Strategy

CVPR 2025 · Architectures and efficiency › RNNs and normalization · approach: Dataset condensation and flow estimation · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Inf-CL, a tile-based, multi-level (cross-GPU ring + in-GPU fused kernel) computation of the contrastive loss log-sum-exp that avoids materializing the full similarity matrix, enabling CLIP training at batch sizes up to millions.

Problem

Contrastive loss memory grows quadratically with batch size because the full image-text similarity matrix must be instantiated for the log-sum-exp, limiting large-batch contrastive training.

Approach

Decompose the contrastive loss into a positive term and a log-sum-exp term, compute LSE over tiles of the similarity matrix with a numerically stable cumulative merge, recompute tiles in the backward pass, and combine ring-based cross-GPU communication with fused in-GPU kernels.

Key results

• Inf-CL computes contrastive loss at batch size 256k on 8 GPUs using 1.63 GB per GPU, a 78x reduction versus prior methods.
• Max batch size on 8×A800 for ViT-L/14 rises from 152k (OpenCLIP) to 448k (Inf-CL), and to 4096k with data offload.
• Zero-shot ImageNet accuracy at 64k batch is 74.93 for Inf-CL vs 74.74 for vanilla CLIP, indicating no accuracy loss.

Measured effects reported in the paper

Surprising findings

• Excessively large batch size (1024k) yields suboptimal zero-shot performance under the authors' settings.
• Tiled implementation is comparable or faster in speed than prior methods despite serial tile processing, attributed to reduced HBM I/O and loss being a small fraction of iteration time.

Limitations

• Peak memory still grows with batch size due to data memory, requiring a data offload workaround.
• Very large batch sizes (1024k) degraded zero-shot performance under the tested hyperparameters.
• Accuracy evaluation mostly uses frozen pretrained image encoder (LiT-style) with ViT-B models, not full end-to-end training at largest scales.

Experimental setup

Datasets: LAION-400M CC3M CC12M

Benchmarks: ImageNet ImageNet-v2 ObjectNet ImageNet-OOD MSCOCO

Metrics: zero-shot top-1 accuracy peak GPU memory (GB) maximum batch size training speed (time per iteration/epoch) image-to-text and text-to-image retrieval R@1

Baselines: CLIP vanilla loss OpenCLIP local loss DisCo-CLIP Gradient Cache

Builds on: CLIP Gradient Cache FlashAttention Ring Attention OpenCLIP DisCo-CLIP

Closest papers by meaning

• On the Similarities of Embeddings in Contrastive Learning (ICML 2025 · similarity 0.83)
• $\mathbb{X}$-Sample Contrastive Loss: Improving Contrastive Learning with Sample Similarity Graphs (ICLR 2025 · similarity 0.83)
• Bridging Mini-Batch and Asymptotic Analysis in Contrastive Learning: From InfoNCE to Kernel-Based Losses (ICML 2024 · similarity 0.83)
• BECLR: Batch Enhanced Contrastive Few-Shot Learning (ICLR 2024 · similarity 0.83)
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (ICLR 2024 · similarity 0.82)
• Communication-Efficient LLM Adaptation over Decentralized GPU Meshes (NeurIPS 2026 · similarity 0.82)