Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss

## Abstract

Abstract Contrastive loss is a powerful approach for representation learning, where larger batch sizes enhance performance by providing more negative samples to better distinguish between similar and dissimilar data. However, scaling batch sizes is constrained by the quadratic growth in GPU memory consumption, primarily due to the full instantiation of the similarity matrix.
To address this, we propose a tile-based computation strategy that partitions the contrastive loss calculation to arbitrary small blocks, avoiding full materialization of the similarity matrix. Furthermore, we introduce a multi-level tiling strategy to leverage the hierarchical structure of distributed systems, employing ring-based communication at the GPU level to optimize synchronization and fused kernels at the CUDA core level to reduce I/O overhead.
Experimental results show that the proposed method scales batch sizes to unprecedented levels. For instance, it enables contrastive training of a CLIP-ViT-L/14 model with a batch size of 4M or 12M using 8 or 32 A800 80GB without sacrificing any accuracy.
Compared to SOTA memory-efficient solutions, it achieves a two-order-of-magnitude reduction in memory while maintaining comparable speed. The code will be made publicly available.

## 1 Introduction

Contrastive learning serves as a foundational technique across various applications, such as multi-modality retrieval (Radford et al., ; Luo et al., ; Girdhar et al., ), self-supervised representation learning (Chen et al., ; He et al., ; Gao et al., ), and dense text retrieval (Wang et al., ).
It learns an embedding space in which similar data pairs stay close while dissimilar ones are far apart (Hadsell et al., ; Oord et al., ; Weng, ).
Large batch sizes are critical to the success of contrastive learning due to their reliance on in-batch negatives (Chen et al., ; Radford et al., ). Specifically, larger batches provide a diverse set of negative samples, enhancing the model’s ability to learn discriminative representations (Pham et al., ).

Despite the above benefits, scaling batch size in contrastive learning is severely limited by GPU memory. The memory needed for computing and storing image-text similarity matrices (Figure 2(a)) grows quadratically with batch size, making further scaling impractical and limiting potential performance gains, even with advanced hardware.
Several methods have been proposed to mitigate memory limitations when scaling batch sizes in contrastive learning.
Gradient-Cache (Gao et al., ) reduces memory usage by decoupling model and loss computations, but the memory cost of the loss still poses a significant bottleneck.
OpenCLIP (Ilharco et al., ) and DisCo-CLIP (Chen et al., ) enhance efficiency by distributing contrastive loss computation across $n$ GPUs, reducing memory consumption by a factor of $n$.
Despite advances in memory-efficient techniques, most studies are limited to a batch size of 128$k$, restricting the potential of contrastive learning and the scaling demands of modern models and datasets (Chen et al., ; Kaplan et al., ).

Figure: Figure 2: (a) Vanilla implementation of contrastive loss gathers features to all devices to calculate all similarity simultaneously, where the similarity with squared complexity are repeatedly stored in all devices, causing huge memory costs for loss calculation when batch size increases. (b) Our Inf-CL significant decreases the memory cost by serial and distributed tile-wise computation.
Refer to caption: x2.png

In this paper, we introduce Inf-CL, a novel approach to mitigate the quadratic memory cost in contrastive learning, which is caused by the full instantiation of the similarity matrix for log-sum-exp (LSE) computation. Instead of storing the entire matrix, Inf-CL partitions the LSE calculation into smaller, sequentially computed tiles, leveraging the cumulative property of LSE. This confines memory usage to the tile size and the number of parallel tiles, allowing for a trade-off between memory and computational efficiency.
To enhance practical efficiency, we propose a multi-level tiling strategy. At a coarse-grained level, image and text batches are distributed across multiple GPUs, with each GPU performing serial LSE computations on multiple rows. As computations proceed, asynchronous column-wise data exchange minimizes communication overhead, as illustrated in Figure 2(b). At a fine-grained level, row-wise computations are parallelized across CUDA cores within each GPU, consolidating iterations into a single kernel to reduce I/O overhead.
Theoretically, Inf-CL can compute contrastive loss with nearly infinite batch sizes using a small tile size, albeit with reduced speed. The multi-level tiling strategy is crucial to achieving practical scalability and efficiency, balancing memory reduction with computation speed.

We evaluate Inf-CL on the image-text contrastive learning task. As shown in Figure 1, Inf-CL reduces space complexity from quadratic (e.g., $\mathcal{O}(b^{2})$ for CLIP, $\mathcal{O}(b^{2}/n)$ for OpenCLIP) to linear ($\mathcal{O}(b/n^{2})$ for Inf-CL), where $b$ and $n$ are the batch size and the number of GPUs.
This substantial reduction in memory usage allows efficient training with large batch sizes.
For instance, training a ViT-L/14 CLIP model with a batch size over 10M on 32 A800 GPUs (80 GB each) requires only 1.44 GB of memory per GPU—over a 30× improvement over previous methods. Moreover, Inf-CL maintains precision consistent with existing approaches. In terms of computation time, Inf-CL matches the performance of prior methods, taking approximately 59 hours to process a 64k batch size on 8 A800 GPUs. The time cost scales nearly linearly with batch size, as demonstrated by a batch size increase from 64k to 256k resulting in a roughly 4× growth in training time ($220.3/49.4\approx 4$).

In summary, our contributions include:

- •
We propose a tile-based contrastive loss implementation that iteratively accumulates the LSE term, removing the need to instantiate the full similarity matrix and significantly reducing memory overhead. This approach theoretically allows training with nearly infinite batch sizes using sufficiently small tiles.
- •
We propose a multi-level tiling strategy for a distributed training system, which reasonably leverages parallelism to achieve a balance between memory and computational efficiency.
- •
Our experiments demonstrate that Inf-CL scales batch sizes to unprecedented levels (e.g., 12M for CLIP-ViT-L/14 on 32 A800 80GB GPUs) while maintaining accuracy and comparable training speed to state-of-the-art methods.

## 2 Preliminaries

### 2.1 Distributed training system

Cross-GPU Communication:
For scaling batch size, training across multiple GPUs is crucial to handle memory and computational demands. However, communication overhead between GPUs can limit performance. Techniques like hierarchical all-reduce and ring-based communication alleviate such overhead by optimizing synchronization between GPUs (Liu et al., ). Blockwise parallelism, as employed in methods like ring attention, further improves efficiency by overlapping computation and communication.

GPU Memory and Execution:
The performance of modern deep learning models relies heavily on hardware resources, particularly GPU memory and execution capabilities.
GPUs, like A100s, typically have two different types of memory: HBM (High Bandwidth Memory) and SRAM (Static Random Access Memory).
HBM serves as the primary memory with a capacity of up to 80GB.
In contrast, SRAM is much smaller (usually measured in megabytes) but offers a significantly faster access speed, acting as a vital cache for frequently accessed data and enabling rapid computations.
Techniques like FlashAttention (Dao et al., ) show that fine-grained control over the memory access of HBM and the fuse the operations can achieve faster training and less memory usage.

### 2.2 Vanilla Implementation of Contrastive Loss

In contrastive learning, the objective is to learn an embedding space where similar samples (positive pairs) are pulled closer, while dissimilar samples (negative pairs) are pushed away.
A typical implementation, exemplified by CLIP (Radford et al., ), is depicted in Figure 2. The image and text encoders are trained with contrastive loss after extracting features. For brevity, we only discuss image-to-text contrastive loss as an example in the following sections, since the implementation of text-to-image loss is symmetric.
Specifically, given a batch size of $b$, the in-batch $c$-dimensional visual feature ${\bm{I}}\in{\mathbb{R}}^{b\times c}$, and textual feature ${\bm{T}}\in{\mathbb{R}}^{b\times c}$, the contrastive loss is defined as

$$ $\mathcal{L}_{I}=-\frac{1}{b}\sum_{i=1}^{b}\log\frac{e^{x_{i,i}}}{\sum_{j=1}^{b }e^{x_{i,j}}},$ (1) $$

where $x_{i,j}={\bm{I}}_{i}\cdot{\bm{T}}_{j}$ is the scaled cosine similarity between the $i$-th image and $j$-th text, and $x_{i,i}$ represents the positive pair. Here, we omitted the temperature factor for simplicity.

The vanilla implementation first computes the similarity matrix ${\bm{X}}\in{\mathbb{R}}^{b\times b}={\bm{I}}\cdot{\bm{T}}^{\prime}$ and stores it in high-bandwidth memory (HBM). Afterward, softmax normalization followed by the calculation of negative log-likelihood is applied to the similarity matrix. The memory required to store ${\bm{X}}$ and its normalized results scales as $\mathcal{O}(b^{2})$, which can occupy a substantial amount of GPU memory when $b$ is large.
Figure 2 gives an example of training ViT-B/16 with a batch size of 64k, using model memory optimization techniques such as Gradient Cache (Gao et al., ; Pham et al., ). As can be seen, the GPU memory footprint of the model itself is only 5.24GB while the loss calculation still requires 66GB.
This indicates that, with batch size scaling, the memory bottleneck during training lies in the loss calculation. Although large batch sizes are necessary for improving model performance (Saunshi et al., ; Chen et al., ), the traditional implementation struggles to support them due to excessive memory consumption in the loss calculation.

## 3 Method

### 3.1 Tile-wise Contrastive Learning

As discussed in Section 2.2, the root cause of the quadratic memory growth in the vanilla implementation is the full materialization of the similarity matrix ${\bm{X}}$.
To eliminate the memory cost, we first decompose the operations related to ${\bm{X}}$ from the loss function:

$$ $\mathcal{L}_{I}=-\frac{1}{b}\sum_{i=1}^{b}(x_{i,i}-\log\sum_{j=1}^{b}e^{x_{i,j }})=-\frac{1}{b}\sum_{i=1}^{b}x_{i,i}+\frac{1}{b}\sum_{i=1}^{b}\log\sum_{j=1}^ {b}e^{x_{i,j}},$ (2) $$

where the spatial complexity of the first part is $\mathcal{O}(b)$, and for the second log-sum-exp (LSE) part, it is $\mathcal{O}(b^{2})$.
Based on this formulation, we introduce a tile-wise contrastive learning method that avoids the full materialization of ${\bm{X}}$ by iterative accumulation between tiles.
The following sections provide a detailed formulation of the forward and backward processes.

Tile-Wise Forward.
To reduce the dependency on storing ${\bm{X}}$ entirely, we adopt a tile-wise approach for calculating ${\bm{l}}$. The process is show as below:

$$ $\underbrace{\left[\begin{array}[]{ccc}{\bm{X}}^{1,1}&\cdots&{\bm{X}}^{1,n_{c}} \\ \vdots&\ddots&\vdots\\ {\bm{X}}^{n_{r},1}&\cdots&{\bm{X}}^{n_{r},n_{c}}\end{array}\right]}_{\text{ Tiled computation of }{\bm{X}}}\rightarrow\underbrace{\left[\begin{array}[]{ ccc}{\bm{l}}^{1,1}&\cdots&{\bm{l}}^{1,n_{c}}\\ \vdots&\ddots&\vdots\\ {\bm{l}}^{n_{r},1}&\cdots&{\bm{l}}^{n_{r},n_{c}}\end{array}\right]}_{\text{ Merged serially via Eq.~{}\ref{eq:tile_lse}}}\rightarrow\begin{pmatrix}{\bm{l} }^{1}\\ \vdots\\ {\bm{l}}^{n_{r}}\\ \end{pmatrix}={\bm{l}}$ (3) $$

where $n_{r}$ and $n_{c}$ represent the number of tiles along the rows and columns, respectively. The computation proceeds by dividing ${\bm{X}}$ into multiple tiles, denoted as ${\bm{X}}^{i,j}$, and then calculating the intermediate LSE values ${\bm{l}}^{i,j}=\mathrm{LSE}({\bm{X}}^{i,j})$ within each tile. The resulting LSE values from each column of tiles are then merged serially along the rows to obtain the final global LSE vector ${\bm{l}}$.

To prevent numerical instability and overflow during the merging process, the following numerically stable operation is performed:

$$ ${\bm{l}}^{i}\leftarrow{\bm{l}}^{i}+\log(1+e^{{\bm{l}}^{i,j}-{\bm{l}}^{i}}),\ j =1,\dots,n_{c},$ (4) $$

where the initial value of ${\bm{l}}^{i}$ is 0. In each iteration, the intermediate value ${\bm{l}}^{i,j}$ is merged with ${\bm{l}}^{i}$, and after processing all $n_{c}$ tiles, the global LSE vector ${\bm{l}}$ is obtained.

During the computation of $\mathrm{LSE}({\bm{X}}^{i,j})$, direct exponentiation can lead to numerical overflow. To address this, we compute ${\bm{l}}^{i,j}$ using the following stabilized formulation:

$$ ${\bm{l}}^{i,j}=\log\sum_{k}e^{{\bm{X}}^{i,j}_{:,k}}={\bm{m}}^{i,j}+\log\sum_{k }e^{{\bm{X}}^{i,j}_{:,k}-{\bm{m}}^{i,j}},$ (5) $$

where ${\bm{m}}^{i,j}=\max_{k}{\bm{X}}^{i,j}_{:,k}$ is a vector, with each element representing the maximum value of the corresponding row in ${\bm{X}}^{i,j}$. This vector acts as a normalization factor, ensuring that the values inside the exponential function remain numerically stable.

This tile-wise approach significantly reduces the memory requirement by allowing each GPU to compute and store only a subset of the similarity matrix at any given time, rather than the entire $b\times b$ matrix.
Additionally, this method facilitates scaling to larger batch sizes by enabling parallel computation of the tiles on multiple GPUs or across different nodes in a distributed system.

Tile-Wise Backward.
According to the chain rule, the gradients *w.r.t.* ${\bm{I}}_{i}$ and ${\bm{T}}_{j}$ are

$$ $\frac{\partial\mathcal{L}_{I}}{\partial{\bm{I}}_{i}}=\sum_{j}\frac{\partial \mathcal{L}_{I}}{\partial x_{i,j}}\cdot\frac{\partial x_{i,j}}{\partial{\bm{I} }_{i}},\quad\ \frac{\partial\mathcal{L}_{I}}{\partial{\bm{T}}_{j}}=\sum_{i} \frac{\partial\mathcal{L}_{I}}{\partial x_{i,j}}\cdot\frac{\partial x_{i,j}}{ \partial{\bm{T}}_{j}}.$ (6) $$

Taking the gradients *w.r.t.* ${\bm{I}}_{i}$ as an example, according to Equation 2, the complete formulation is

$$ $\displaystyle\frac{\partial\mathcal{L}_{I}}{\partial{\bm{I}}_{i}}$ $\displaystyle=-\frac{1}{b}\sum_{j}(\frac{\partial\mathcal{L}_{I}}{\partial x_{ i,i}}\cdot\frac{\partial x_{i,i}}{\partial x_{i,j}}\cdot\frac{\partial x_{i,j} }{\partial{\bm{I}}_{i}}-\frac{\partial\mathcal{L}_{I}}{\partial l_{i}}\cdot \frac{\partial l_{i}}{\partial x_{i,j}}\cdot\frac{\partial x_{i,j}}{\partial{ \bm{I}}_{i}})$ (7) $\displaystyle=-\frac{1}{b}\cdot{\bm{T}}_{i}+\frac{1}{b}\sum_{j}e^{x_{i,j}-l_{i }}\cdot{\bm{T}}_{j}.$ $$

From the formula, it can be seen that the second term requires the similarities $x_{i,j}$ with $\mathcal{O}(b^{2})$ memory in common implementations, whether stored in the forward process or computed directly in the backward process.
To tackle this, we apply the similar tile-based method as the forward process to compute the gradient.
Specifically, we first store ${\bm{l}}$, which has only $b$ elements during forward propagation, and calculate the gradient *w.r.t* ${\bm{I}}_{i}$ by iterative accumulation in multiple tiles:

$$ $\displaystyle{\bm{I}}_{i}^{\prime}\leftarrow{\bm{I}}_{i}^{\prime}+e^{x_{i,j}-l _{i}}\cdot{\bm{T}}_{j},\ j=1,\dots,n_{c},$ (8) $\displaystyle\frac{\partial\mathcal{L}_{I}}{\partial{\bm{I}}_{i}}=-\frac{1}{b} \cdot{\bm{T}}_{i}+\frac{1}{b}{\bm{I}}_{i}^{\prime},$ $$

where ${\bm{I}}_{i}^{\prime}$ is a temporary variable for accumulation.
The detailed algorithm is shown in Appendix.

### 3.2 Multi-Level Tiling

Figure: Figure 3: Multi-level tiling strategy. Top: for cross-GPU tiling, each GPU is assigned with multiple rows. The computation and the column-wise communication are performed asynchronously to reduce the cost. Bottom: for in-GPU tiling, the calculations in each GPU are further divided into tiles and the row-wise calculation is distributed to multiple CUDA cores. The accumulative operations of each row are merged into one kernel for reducing I/O times between SRAM and HBM.
Refer to caption: x3.png

The scaling of batch size is usually accompanied by the scaling of the number of GPUs.
In order to fully utilize the parallelism between multiple GPUs while exploiting partially serial computation on a single GPU to reduce the memory cost, we propose a multi-level tiling method that distributes the above LSE calculation to coarse-grained cross-GPU tiles and fine-grained in-GPU tiles.

Cross-GPU Tile.
As shown in Algorithm 1, in data parallel training with $n$ GPUs, the $i$-th GPU first processes a portion of images and texts to visual features ${\bm{I}}^{i}\in{\mathbb{R}}^{b_{s}\times c}$ and textual features ${\bm{T}}^{i}\in{\mathbb{R}}^{b_{s}\times c}$, where $b_{s}=b/n$ is the batch size in one GPU.
Then for the calculation of the contrastive loss, we distribute computations of different rows to different GPUs and synchronize the columns between GPUs step-by-step, considering the row-wise characteristic.
Specifically, the $i$-th GPU is responsible for calculating ${\bm{X}}^{i,:}$ and the corresponding ${\bm{l}}^{i}$.
For memory considerations, based on the tiling strategy described in Section 3.1 where only one tile ${\bm{X}}^{i,j}$ is computed at a time, ${\bm{X}}^{i,:}$ is further divided into ${\bm{X}}^{i,j}$ for $n$ step to calculate ${\bm{l}}^{i}$ following Equation 4, where the local LSE ${\bm{l}}^{i,j}$ is calculated by in-gpu tiling as described in the next part.

Moreover, since the computation of ${\bm{X}}^{i,j}$ while $i\neq j$ requires the textual feature ${\bm{T}}^{j}$ stored in other GPUs, additional communication overhead is inevitable, especially as the number of GPUs grows.
In order to reduce or even eliminate the communication overhead, we associate all GPUs with a ring topology, based on the idea of overlapping communication time and computation time overlap as much as possible.
Concretely, starting with ${\bm{T}}^{i}$, each GPU process sends the current textual features to the next process and receives the textual features from the previous process using the ring topology while computing Equation 4.
In this way, the communication time cost is negligible when it is greater than the computation time overhead.

Figure: Algorithm 1 Forward Process of Multi-level Tile-Wise Global LSE Calculation

In-GPU Tile.
With the cross-GPU tiling technique, the memory complexity becomes $\mathcal{O}(b_{s}^{2})$ for directly storing ${\bm{X}}^{i,j}$ where $b_{s}=b/n$.
Since the number of GPU $n$ is somehow limited, we further introduce in-GPU tiling to reduce the $\mathcal{O}(b_{s}^{2})$ memory cost to $\mathcal{O}(b_{s})$ for enabling further batch size scaling.
Specifically, we first split $\tilde{{\bm{X}}}={\bm{X}}^{i,j}$ into tiles:

$$ $\tilde{{\bm{X}}}=[\tilde{{\bm{X}}}^{i,j}],\ i=1,\dots,\tilde{n}_{r},\ j=1, \dots,\tilde{n}_{c},$ (9) $$

where $\tilde{n}_{r}=\lceil b/t_{r}\rceil$ and $\tilde{n}_{c}=\lceil b/t_{c}\rceil$ and $t_{r}$ and $t_{c}$ is the row-wise and column-wise size of a tile.
For implementation, we distribute rows to multiple CUDA cores to make full use of the parallel computing power of the GPU, and serial process the row-wise tiles in each kernel by applying Equation 5 and Equation 4 to $\tilde{{\bm{X}}}^{i,j}$, as shown in Algorithm 2.

The iterative computation requires multiple memory access for variable ${\bm{l}}^{i}$.
To avoid expensive I/O from HBM to SRAM, we fuse the row-wise iterative calculation into one kernel.
Specifically, ${\bm{l}}^{i}$ and $\tilde{{\bm{X}}}^{i,j}$ are allocated in SRAM.
In this way, the image features are loaded to SRAM only once at beginning, and $\tilde{{\bm{l}}}^{i}$ is written to HBM only once in the end, as shown in Figure 3.

Figure: Algorithm 2 Forward Process of Tile-Wise Local LSE Calculation

**Table 1: Training Memory Cost Across Different Hardware and Batch Sizes. Experiments utilize Data Parallelism with Automatic Mixed Precision for efficient distributed training. The baselines include the Vanilla loss (CLIP) and Local loss (OpenCLIP). To minimize memory consumption, Gradient Cache is adopted, with an accumulation batch size of 128. ^∗ indicates the use of the data offload strategy, which reduces memory usage by transferring only a small data batch from CPU to GPU during each accumulation step. ✗ denotes cases where the baseline exceeds the hardware memory limit for a given batch size, making training infeasible. Memory cost is evaluated using the ViT-L/14 architecture and the AdamW optimizer.**
|  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Model | Loss (Peak) Memory Cost (GB) |  |  |  |  |
|  | 32k | 64k | 128k | 256k | 1024k |
| 8$\times$A800 ($\approx 8\times 80$GB) |  |  |  |  |  |
| CLIP | 16.67 (46.40) | 66.11 (77.94) | ✗ | ✗ | ✗ |
| OpenCLIP | 2.27 (43.97) | 8.63 (46.38) | 33.64 (51.23) | ✗ | ✗ |
| \cdashline1-6<br>Inf-CL | 0.18 (44.20) | 0.36 (46.63) | 0.72 (51.46) | 1.45 (61.13) | ✗ |
| Inf-CL^∗ | 0.18 (42.40) | 0.36 (42.49) | 0.72 (42.69) | 1.45 (43.07) | 6.53 (45.40) |
| 32$\times$A800 ($\approx 32\times$80GB) |  |  |  |  |  |
| CLIP | 16.66 (42.85) | 66.11 (75.52) | ✗ | ✗ | ✗ |
| OpenCLIP | 0.71 (42.46) | 2.45 (43.06) | 8.98 (44.26) | 34.35 (46.71) | ✗ |
| \cdashline1-6<br>Inf-CL | 0.05 (42.48) | 0.09 (43.08) | 0.18 (44.30) | 0.35 (46.71) | 1.44 (61.20) |
|  |  |  |  |  |  |

## 4 Experiments

### 4.1 Experimental Settings

Dataset and Data Processing.
We assess the effectiveness of our Inf-CL on Laion400M dataset (Schuhmann et al., ) where we used 280M (out of 400M) samples for training due to the unavailability of images in the remaining samples.
Images undergo preprocessing using RandomResizedCrop with a crop ratio of $[0.75,1.33]$ and a scale of $[0.08,1.0]$.

Training Hyperparameters. A modified AdaFactor optimizer (Shazeer & Stern, ) is employed for training, following the settings of ViT-g (Zhai et al., ). The optimizer is configured with a learning rate of $1\times 10^{-3}$, weight decay of $1\times 10^{-4}$, and coefficients $\beta_{1}=0.9$ and $\beta_{2}=0.95$ (Zhai et al., ). Training spans 8 epochs, using a cosine learning rate schedule with a linear warm-up during the first 0.5 epoch.

Implementation Details. For distributed training, we employ Data Parallelism (Li et al., ) with Automatic Mixed Precision (float16)(Micikevicius et al., ). To support larger batch sizes, we adopt Gradient Cache (Gao et al., ) which decouples contrastive loss computation from the model’s forward and backward passes. Consequently, the peak memory cost per iteration, $M_{peak}$, is calculated as:

$$ $M_{peak}\approx M_{data}+\max(M_{loss},M_{backbone}),$ (10) $$

where $M_{data}$ is the memory for data, $M_{loss}$ is for loss computation, and $M_{backbone}$ is for the model’s forward and backward operations.

Baselines. We compare our method against two baselines: the vanilla loss from CLIP and the local loss from OpenCLIP/DisCo-CLIP. The vanilla loss computes a $b\times b$ similarity matrix by gathering both row and column features from all GPUs, while the local loss requires only column features to calculate a $b/n\times b$ similarity matrix, where $b$ and $n$ are the batch size and the number of GPUs.

**Table 2: Maximum batch size for model training using different hardware and contrastive loss methods. The training setting of this experiment is aligned with Table 1.**
|  |  |  |  |  |
| --- | --- | --- | --- | --- |
| Budget | Maximum Batch Size (Loss Memory Cost) | Improvement |  |  |
|  | CLIP | OpenCLIP | Inf-CL | (Ours / Sota) |
| ViT-B/16 |  |  |  |  |
| 8$\times$A800 | 68k (74.39 GB) | 172k (59.95 GB) | 800k  (3.01 GB) | 4.65 (800k/172k) |
| 32$\times$A800 | 68k (74.39 GB) | 360k (66.29 GB) | 3456k (3.27 GB) | 9.60 (3456k/360k) |
| ViT-L/14 |  |  |  |  |
| 8$\times$A800 | 64k (66.11 GB) | 152k (47.23 GB) | 448k  (2.52 GB) | 2.94 (448k/152k) |
| 32$\times$A800 | 64k (66.11 GB) | 352k (64.13 GB) | 2048k (2.89 GB) | 5.82 (2048k/256k) |
| ViT-L/14 w/ data offload |  |  |  |  |
| 8$\times$A800 | 64k (66.11 GB) | 184k (69.10 GB) | 4096k (26.12 GB) | 22.26 (4096k/184k) |
| 32$\times$A800 | 64k (66.11 GB) | 368k (64.13 GB) | 12288k (19.59 GB) | 33.39 (12288k/368k) |
|  |  |  |  |  |

Figure: Figure 4: Training Speed of ViT-L/14 CLIP on 8$\times$A800 for Varying Batch Sizes. The left figure shows the time per iteration step, while the right displays the time per epoch. Loss calculation contributes minimally to the total iteration time, making Inf-CL’s iteration time comparable to previous methods. Furthermore, the iteration time of Inf-CL scales linearly with batch size, leading to a stable training duration of approximately 59 hours per epoch.
Refer to caption: x4.png

### 4.2 Cost Analysis

Our method, as detailed in Section 3.2, divides the calculation of contrastive loss into tiles and distributes them across different GPUs and GPU kernels. To rigorously assess its memory efficiency, we compare our approach with previous methods like CLIP and OpenCLIP by evaluating “Memory Consumption”,“Max Supported Batch Size” and “Speed” across various model architectures and hardware settings. The effective memory cost is determined by peak memory (Equation 10), which is the maximum memory needed during an iteration.

Memory Consumption.
To illustrate the memory efficiency of Inf-CL, we compared it to previous methods using the same batch size. Table 1 shows that for loss calculation, Inf-CL requires significantly less memory than its predecessors. Specifically, with a batch size of 128k on 8$\times$A800, Inf-CL only consumes 0.72 GB, whereas OpenCLIP requires 33.64 GB.
However, while the memory cost of loss calculation with Inf-CL is minimal, peak memory usage still increases rapidly with batch size due to growing data memory, as discussed in “Max Supported Batch Size.”
By integrating Inf-CL with “data offload”, we can mitigate this memory increase, enabling us to train a ViT-L/14 model with a batch size of 1024k on 8$\times$A800.

Maximum Batch Size.
We compare the maximum batch size of Inf-CL with those of previous approaches under various model architectures (ViT-B/16 or ViT-L/14) and training budgets (8$\times$A800 or 32$\times$A800). As shown in Table 2.
Inf-CL significantly outperforms previous SOTA methods, achieving improvements of 4.65$\times$ for ViT-B/16 on 8×A800, which is further increased to 9.60$\times$ when using 32$\times$A800. Notably, as we scale up the model size, the improvements decrease; for instance, from 4.65 to 2.94 when changing from ViT-B/16 to ViT-L/14. To understand this trend, we analyze peak memory usage. Since Inf-CL has negligible memory requirements, peak memory is primarily driven by $M_{backbone}+M_{data}$. $M_{backbone}$ is constant, meaning the rapid growth in peak memory is mainly due to increased $M_{data}$. Since ViT-L/14 has a larger $M_{backbone}$, the remaining memory can accommodate only a smaller batch size for $M_{data}$.
To address this issue, we implement “data offload”, which allows us to load only a small batch of data onto the GPU for each accumulation step, effectively stabilizing the data memory usage.
Therefore, by combining data offload with our Inf-CL, we can scale the batch size to over 10M on 32$\times$A800.

Training Speed.
We compare the training speed of our Inf-CL with previous methods. As shown in Figure 4, using Inf-CL to train ViT-L/14 on 8$\times$A800 has almost the same speed as previous methods.
Even when increasing batch size beyond the limits of previous methods, Inf-CL maintains a linear increase in iteration time, with one epoch consistently taking about 59 hours.
Combining training speed results with memory cost results demonstrates that our Inf-CL has superior memory efficiency, while only introducing a little additional time cost (extra analysis in Appendix A.2).

**Table 3: Performance Verification. The training strategies is consistent with Table 2. We choose ViT-B/16 as the model architecture and adopt LiT strategy like Table 4. We evaluate zero-shot top-1 classification accuracy on several data sets, e.g., ImageNet-Validation Deng et al. (), ImageNet-v2 (Recht et al., ), ObjectNet (Barbu et al., ) and ImageNet-OOD (Hendrycks et al., ). We also evaluate zero-shot image-text top-1 retrieval accuracy on MSCOCO (Chen et al., ).**
|  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Method (Batch Size) | ImageNet | MSCOCO R@1 |  |  |  |  |
|  | Validation | v2 | ObjectNet | OOD | I$\rightarrow$T | T$\rightarrow$I |
| Vanilla     (64K) | 74.74 | 65.30 | 46.31 | 66.13 | 25.71 | 44.31 |
| OpenCLIP     (64K) | 74.86 | 65.22 | 46.29 | 66.75 | 25.98 | 44.02 |
| \cdashline1-6<br>Inf-CL     (64K) | 74.93 | 65.27 | 46.13 | 66.77 | 26.01 | 43.95 |
| Inf-CL   (256K) | 75.12 | 65.12 | 46.44 | 67.15 | 25.90 | 44.61 |
| Inf-CL (1024K) | 73.58 | 63.87 | 44.55 | 64.60 | 24.53 | 41.58 |
|  |  |  |  |  |  |  |

**Table 4: Ablation Study of Multi-level Tiling Strategy. The training strategies is consistent with Table 2, using the ViT-B/16 architecture. To reduce memory consumption and expedite experimentation, we freeze the image encoder and load pretrained weights as done in LiT. The global batch size is fixed at 64k with an accumulation batch size of 256 per GPU. These experiments are conducted on 4$\times$A800 (80G) GPUs. “Complexity” denotes the space complexity of loss calculation. $b$ denotes batch size, while $n$ denotes the number of GPUs.**
|  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Cross-GPU | In-GPU | Data | Loss | Backbone | Peak | ImageNet |  |
|  |  | Memory | Complexity | Memory | Memory | Memory |  |
| (Vanilla) | 1.96 | $\mathcal{O}(b^{2})$ | 66.21 | 8.26 | 69.24 | 74.82 |  |
| (OpenCLIP) | 1.96 | $\mathcal{O}(b^{2}/n)$ | 16.96 | 8.26 | 20.79 | 74.86 |  |
| \cdashline1-8<br>✔ |  | 1.96 | $\mathcal{O}(b^{2}/n^{2})$ | 4.81 | 8.26 | 12.30 | 74.78 |
| ✔ | ✔ | 1.96 | $\mathcal{O}(b/n^{2})$ | 0.81 | 8.26 | 12.30 | 74.93 |
|  |  |  |  |  |  |  |  |

### 4.3 Performance Analysis

In this section, we investigate whether introducing Inf-CL negatively affects CLIP performance and whether increasing batch size with Inf-CL enhances performance.
Due to the limit of GPU resources, we utilize the ViT-B/16 with Bert-Base (Devlin, ). We follow the training strategy of LiT (Zhai et al., ) to freeze the visual backbone and use the pre-trained weights instead.

Performance Verification.
We evaluate CLIP models trained with different loss implementations, with the results presented in Table 3. As shown, under the same batch size, our Inf-CL performs similarly to previous methods, with performance differences falling within the error margin, confirming that our design incurs no precision loss in the loss calculations.
Furthermore, the results indicate that increasing the batch size within a certain range yields performance enhancements, thereby underscoring the significance of our method for helping scale the batch size.
However, under our experimental conditions, we currently observe that an excessively large batch size—previously unexamined in the literatures—results in suboptimal performance.
This may be attributed to factors such as unoptimized hyperparameters, inadequate training iterations, or constraints related to data size (for a comprehensive analysis, see Appendix A.3).
Since our work mainly focus on how to enable large batch size training, these factors warrant further investigation in future work.

Ablation Study.
We ablate multi-level tiling in Table 4 and show that our designs incur no precision loss in loss calculations.
This allows arbitrary combinations to achieve nearly the same zero-shot classification accuracy (about 74.8% on ImageNet for 64k batch size), while significantly reducing memory costs.
According to the Equation 10, their $M_{peak}$ is decided by $M_{backbone}$ + $M_{data}$ rather than $M_{loss}$ + $M_{data}$ as in prior methods.
For complexity analysis, Cross-GPU tiling is $\mathcal{O}(b^{2}/n^{2})$, resulting in a memory cost that is $1/n$ of OpenCLIP ($16.96/4.81\approx 4$ in Table 4).
Based on it, introducing In-GPU tiling can further reduce memory cost and make the growth of memory cost linear, i.e., $\mathcal{O}(b^{2}/n^{2})\rightarrow\mathcal{O}(b/n^{2})$.

## 5 Related Work

Contrastive Learning: The core idea of contrastive learning is to learn better representations by distinguishing between positive and negative pairs of samples (van den Oord et al., ; Chen et al., ). This approach demonstrates strong effectiveness across diverse tasks, as the nature of the paired samples varies depending on the specific application. In image foundation models, such as SimCLR (Chen et al., ) and MoCo (He et al., ), positive pairs are created by augmenting the same image in different ways. For cross-modal retrieval, as exemplified by CLIP (Radford et al., ) and ALIGN (Jia et al., ), the positive pairs consist of aligned image and text samples. Similarly, for dense text retrieval (Karpukhin et al., ; Wang et al., ; Zhang et al., ), the positive pairs are composed of query and document pairs. Several works improve contrastive learning performance by enhancing dataset quality, modifying the loss function, or refining negative sample selection (Vasu et al., ; Zhai et al., ; Zhang et al., ). Moreover, several studies, both empirical and theoretical, have demonstrated from various perspectives that larger batch sizes contribute to learning better representations (Saunshi et al., ; Chen et al., ). Due to the quadratic growth of memory usage with batch size in classical contrastive loss, most existing studies have stopped scaling their batch sizes to 128k, even when leveraging hundreds of GPUs (Radford et al., ; Jia et al., ; Yang et al., ).

Memory-efficient Training: As deep learning models continue to grow in size and complexity, the demand for computational resources, particularly GPU memory, has increased significantly. Techniques such as Gradient Checkpointing (Sohoni et al., ) recompute activations during backpropagation to save memory at the expense of additional computation. Flash Attention (Dao et al., ) reduces memory overhead by computing attention in blocks without storing large intermediate states. Ring Attention (Liu et al., ) distributes long sequence activations across multiple devices, overlapping computation and communication to train sequences far longer than previous methods. For contrastive learning, GradCache (Gao et al., ) and BASIC (Pham et al., ) introduce a gradient caching technique that decouples backpropagation between contrastive loss and the encoder, which reduces memory usage in the model by accumulating gradients per mini-batch. OpenCLIP (Ilharco et al., ) and DisCo-CLIP (Chen et al., ) reducing memory consumption by distributing the computation of contrastive loss across multiple GPUs.

## 6 Conclusion

This paper addresses the GPU memory bottleneck in scaling batch sizes for contrastive loss. To overcome the quadratic memory consumption resulting from the full instantiation of the similarity matrix, we proposed a tile-based computation strategy that partitions the calculation into smaller blocks, thus avoiding full matrix materialization. Furthermore, we introduced a multi-level tiling strategy that leverages ring-based communication and fused kernels to optimize synchronization and minimize I/O overhead. Our experiments demonstrated that our method scales contrastive loss batch sizes to unprecedented levels without compromising accuracy or training speed. This approach marks a significant advancement in large-scale contrastive learning, shedding light on further developments in areas such as self-supervised learning and dense text retrieval.

## Appendix A Appendix

### A.1 Backward Process

Figure: Algorithm 3 Backward Process of Multi-level Tile-Wise Global LSE Calculation

Figure: Algorithm 4 Backward Process from of intra-GPU Tile-Wise LSE calculation

### A.2 Analysis of Training Speed Efficiency in Inf-CL

Although Inf-CL might be expected to exhibit slower performance because it breaks the loss calculation to small tiles and serially process these tiles, it achieves comparable speed to previous methods, as shown in Figure 4.
This is primarily due to two factors:
(1) Loss calculation represents only a minor fraction of the total iteration time, especially for large models, thereby exerting minimal impact on the overall iteration time.
(2) While Inf-CL has similar computational complexity to standard contrastive loss, its tiling approach could introduce some speed overhead due to reduced parallelism.
However, Inf-CL fuses the operations of similarity matrix calculation and softmax, which in regular contrastive loss require two separate communications between SRAM and HBM.
By merging these into a single communication, Inf-CL effectively reduces I/O time, mitigating the cost of serial tile computation.

### A.3 Factors influencing performance when scaling batch size

While larger batch size is theoretically expected to enhance performance Chen et al. (), our experimental results deviate from this expectation.
To better understand this discrepancy, we analyze the factors that impact performance when scaling up batch size.

Hyperparameters.
Although larger batch sizes provide more diverse negative samples for contrastive learning, potentially improving the embedding space, careful tuning of hyperparameters is necessary to ensure model convergence.
Previous research indicates that when increasing batch size, the learning rate should be scaled proportionally to maintain a consistent parameter update norm throughout training (Goyal, ).
Since a fixed learning rate is used across all experiments, this may have contributed to the reduced performance observed with larger batch sizes.
Moreover, prior studies suggest that large batch sizes require longer training epochs to ensure sufficient parameter updates and avoid suboptimal convergence (Hoffer et al., ).
Overall, the performance gains from larger batch sizes are contingent on the careful tuning of multiple hyperparameters beyond just learning rate and epochs, highlighting the importance of comprehensive hyperparameter optimization to fully exploit the benefits of scaling.

Figure: Figure 5: Performance of ViT-B/32 across Varying Batch Sizes. Except batch size, other experiment settings are consistent. In Figure, the most suitable batch size is increasing with data scale.
Refer to caption: x5.png

Data Scale.
Increasing batch size improves the precision of gradient estimation for the representation distribution defined by the dataset Chen et al. ().
Larger datasets more accurately capture real-world distributions, and thus, employing a larger batch size enables contrastive loss to generate more precise gradients, enhancing the model’s ability to learn discriminative representations.
As shown in Figure 5, our experiments on different data scales (e.g., CC3M, CC12M and Laion400M) indicate that the optimal batch size increases with dataset size.
Specifically, performance on CC12M saturates at a batch size of 32k, whereas Laion400M achieves saturation at a batch size of 256k.

In summary, while scaling up batch sizes is critical for enhancing contrastive learning, our findings suggest that performance does not monotonically improve with batch size increases. As seen in our previous experiments (Table 3), extremely large batch sizes (e.g., 1024k) can lead to a decline in performance, indicating that factors such as hyperparameter tuning and dataset scale are among the many considerations that influence model effectiveness. This highlights the need for a balanced approach when increasing batch sizes, ensuring that optimal configurations are found to fully exploit the benefits of contrastive learning.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1.	Can you provide statistical measures such as the mean and standard deviation for the results.
2.	Why does the paper focus only on image-text contrastive learning task?
soundness: 3
strengths: The paper is well-written, presenting its ideas clearly and systematically. It effectively balances theoretical insights with practical experimentation, demonstrating a strong command of both. The inclusion of detailed figures and tables enhances readability and helps illustrate the method’s advantages.
summary: The authors propose a tile-based computation strategy to avoid full instantiation of the similarity matrix, which is crucial for efficiently handling large batch sizes in contrastive learning without hitting memory limits. This approach seems impactful and demonstrate that Inf-CL achieves significant reduction in memory use compared to baseline.
weaknesses: W1. The method's core motivation appears to stem directly from the tiling strategy detailed in the FlashAttention paper, raising questions about its originality. Furthermore, while the implementation borrows parameters from SigLip, there is no in-depth comparison or critical evaluation of SigLip, leaving its relative contributions unclear. Several other key comparisons like with DisCo-Clip is missing in the experimental section.  

W2. The results show diminishing returns with increasing batch sizes, eventually leading to performance decline, as also observed in the SigLip paper. This raises doubts about the claim of "infinite batch scaling," which seems overstated. Additionally, the reasons provided in the appendix for this decline e.g. unoptimized hyperparameters and limited training, does not sufficiently supported by experimental evidence. 

W3. While the paper emphasizes memory reduction through tiling strategies, it fails to detail the exact computational overhead introduced by tiling. Without reporting the specific time costs associated with tiling, beyond overall training time, it is difficult to assess the true benefits of the proposed method.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Q1. With respect to Lines 114-115, the authors must elaborate on what specific aspects of their multi-level tiling implementation (which exists in literature) they consider novel or uniquely suited to contrastive learning tasks ? IN current presentation, it seems the authors propose multi-level tiling.

Q2. The authors must discuss how their approach relates to or builds upon key works in loop tiling and multi-level tiling from the systems community https://scholar.google.com/scholar?hl=en&as_sdt=0%2C3&q=Multi-level+tiling&btnG= and https://scholar.google.com/scholar?hl=en&as_sdt=0%2C3&q=loop+tiling&btnG=

Q3. Is the proposed approach unique to CL or generally applicable to any task that requires construction of similar matrix with quadratic cost? In that spirit, the optimization techniques like multi-level tiling and ring all-reduce should already be used for distributing computations as core system level optimization.  Please discuss potential applications of the proposed method to other tasks involving quadratic-cost similarity matrices, and to explain any unique challenges or adaptations required for contrastive learning. 

Q4. What is impact of asynchronous communication on the performance ?

Q5. Figure 2: It is challenging to observe the color schemes in tiles and with printed version, the colors wash out. The authors are encouraged to use deeper colors which are print friendly.
soundness: 3
strengths: 1. The work tackles an important problem for Contrastive Loss (CL), i.e., how to efficiently store and compute similarly matrix that is intrinsically quadratic in batch size. The resulting impact is on how well the CL for image-text task could scale with large batch sizes for large CLIP-ViT models without sacrificing any accuracy. The empirical results demonstrate that proposed technique reduced memory cost by 78x for batch size 256k. 

2. The paper is mostly clearly written with some scope of further improvement. 

3. The experiments are extensive and informative supporting the claimed contributions.
summary: The authors propose a technique Inf-CL for scaling the computations of the Contrastive Loss (CL) for Image-Text tasks. The vanilla CL relies on computing a similarity matrix between Image and Text pairs which has quadratic memory cost with respect to batch size. 

With large batch sizes, that is necessary for improving the performance of Image-text CL, such similarity matrix computation becomes a bottleneck. As a result, there have been recent attempts to distribute the construction and computation of similarity matrix and the resulting CL task.

The proposed technique linearizes the memory cost of similarity matrix by employing  several optimizations such as multi-level tiling and ring all-reduce for coordinating the tasks across distributed GPU devices. In addition, the computation cost scales linearly with batch size with fine-grained tiling with smaller tiles.
weaknesses: 1. The techniques used multi-level tiling and ring all-reduce are well known optimization tricks in hardware and system communities. the authors employ it in novel setup of CL with Image-Text pairs. In principle, one can argue that any task (not necessarily CL) that requires construction of a symmetric similarity matrix having quadratic memory cost would employ such strategies at core system level - be it distributing across devices or across internal cores/threads. For instance, in ViT models, $QK^T$ is a similarity matrix between queries and keys that is quadratic is number of tokens. Moreover, softmax() is also applied row-wise similar to loss computation in CL task (Equation 1 in paper). As a consequence, such optimization strategies would be standard treatment in parallelizing such tasks. 

2. Tiling is a prevalent loop transformation technique that enhances both parallelism and data locality. High-performance implementations often utilize multiple levels of tiling to take full advantage of the parallelism hierarchy and optimize cache and register efficiency. Please refer to this missing reference "Multi-level tiling: M for the price of one" https://dl.acm.org/doi/abs/10.1145/1362622.1362691 that talk about multi-level tiling. As a result, I am not very convinced about the claims that Multi-level tiling is a novel contribution in the manuscript where authors state: "We propose a multi-level tiling strategy for a distributed training system, which reasonably leverages parallelism to achieve a balance between memory and computational efficiency." (Line 114-115) and Section 3.2 dedicated to Multi-Level tiling.  An extensive list of references from systems community is missing that talks about loop tiling and multi-level tiling.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1.	Is there a clear benefit of using large batches in terms of accuracy and training time?
2.	Figure 5 in the appendix indicates that Laion400M reaches saturation at a batch size of 256k. What is the justification for using a batch size greater than 256k for this dataset?
3.	Can we expect to see similar performance with other models and datasets? What level of performance should we anticipate when utilizing thousands of GPUs?
4.	Although the authors mention that larger batches provide a diverse set of negative samples,  negative sampling was not used in the experiments? Does negative sampling have an impact on accuracy?
soundness: 3
strengths: 1.	The paper clearly demonstrates that tiling effectively reduces the memory requirements for contrastive language-image learning. When large batches are desired, the proposed approach offers a practical and scalable solution to this memory challenge.
2.	Inf-CL explores various approaches to prevent numerical overflow. 
3.	The multi-level algorithm uses a ring communication algorithm to share features among GPUs. This approach enables overlapping the communication with computations.
summary: Contrastive language-image learning can require O(b²) memory for a batch size of b, which includes both image and text features. This quadratic memory requirement may pose challenges for large batches. The paper addresses this limitation through a tile-based computation strategy that partitions the contrastive loss calculation into smaller blocks. Additionally, a hierarchical blocking structure has been developed for multi-GPU systems. Experiments have demonstrated that this proposed blocking approach significantly reduces memory requirements while maintaining comparable accuracy and training time.
weaknesses: 1. The paper does not clearly establish the necessity for large batches. The experimental results presented indicate that larger batch sizes do not lead to improvements in either accuracy or runtime. For instance, Figure 5 in the appendix shows that Laion400M reaches saturation at a batch size of 256k. This raises the question of why experiments were conducted with a batch size of 12,288k, as shown in Table 2. Additionally, Cherti et al. demonstrated (in Table 12 of their paper) that performance variations due to changes in batch size are minimal, typically in the range of 0.2% to 0.5% across different settings. Consequently, it remains unclear how the proposed algorithm would be beneficial for contrastive learning.
Ref: Cherti, Mehdi, Romain Beaumont, Ross Wightman, Mitchell Wortsman, Gabriel Ilharco, Cade Gordon, Christoph Schuhmann, Ludwig Schmidt, and Jenia Jitsev. "Reproducible scaling laws for contrastive language-image learning." In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2818-2829. 2023. 

2. All experiments in the paper were conducted using only one dataset (Laion404M), making it unclear whether the results would hold for larger datasets referenced in the OpenCLIP benchmark. 

3. The paper primarily focuses on two models, ViT-L/14 and ViT-B/16, while the OpenCLIP benchmark indicates that other models yield better results. It might be useful to clarify why the reason behind selecting ViT-L/14 and ViT-B/16 for this paper.

4. The title of the paper may be misleading, as it does not address general contrastive loss but focuses on contrastive language-image learning. It might be helpful for clarity if the title reflects the specific contributions of the paper.

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 2
questions: - While the memory footprint reductions are impressive, how meaningful are large to huge batch sizes for models like (Open-)CLIP? Would a large batch size practically not often lead to worse out-of-sample performance?

- The authors correctly point out the symmetry of the contrastive loss, but do not meaningfully exploit this property. If the loss is symmetric, so is the resulting distance matrix. As a result the upper, respectively lower triangle of the matrix are identical to one another. Hence, the computation scheme could be further improved by only calculating on of the two and replicate the other. Have the authors considered this? If so, why is the presented algorithm not implemented as such?
soundness: 3
strengths: - Clear and concise writing, logical structure of the manuscript

- Meaningful visual items to explain the algorithm's principles and experimental results

- More than solid improvements on the memory footprint for computing the contrastive loss

- No loss in accuracy of the algorithm

- Intention to release the source code for reproducibility purposes
summary: The authors introduce Inf-CL, a way to drastically improve the (parallel) computation of contrastive losses and to scale this computation to massive batch sizes. The core issue with the current way of computing the constrastive loss is its quadratic memory consumption that is completely materialized in-memory. In this work, the authors instead propose a block-wise tiled computation of this quadratic matrix. As a result, as long as the block fits into memory, the entire contrastive loss matrix may be computed. Block sizes can be chosen arbitrarily small. Furthermore, they propose a ring communication scheme for distributed-memory system, e.g. a multi-GPU setups. In an experimental evaluation on typical benchmark datasets and models they achieve identical predictive performance, while being able to drastically increase the batch size/size of the constrastive loss matrix.
weaknesses: - It would be meaningful to have the first figure after the abstract. The authors are kindly requested to rearrange both textual elements

- The manuscript's claim of "near-infinite" scaling seems overly sensationalist. Could the authors possibly soften the wording?

- While not conceptually wrong, Algorithm 1 requires several inputs that are not used in pseudo code. It would seem meaningful to explicitly show the use of the required inputs.

- Violation of the conference template, e.g.
    * Caption of tables are above
    * Table content should not be an image but selectable text

- Computing a (symmetric) distance matrix (here: a loss, but algorithmically identical) in parallel on distributed-memory system is a well-studied problem. In principle the communication scheme boils down to a merry-go-round principle, even applied to other matrix computation problems. It would be meaningful to add at least the following works and credit them:
    * Bischof, C. H. (1989). Computing the singular value decomposition on a distributed system of vector processors. Parallel Computing, 11(2), 171-186.
    * Man, D., Uda, K., Ueyama, H., Ito, Y., & Nakano, K. (2010, November). Implementations of parallel computation of Euclidean distance map in multicore processors and GPUs. In 2010 First International Conference on Networking and Computing (pp. 120-127). IEEE.
    * Al-Neama, M. W., Reda, N. M., & Ghaleb, F. F. (2014). An improved distance matrix computation algorithm for multicore clusters. BioMed research international, 2014(1), 406178.
    * Angeletti, M., Bonny, J. M., & Koko, J. (2019). Parallel euclidean distance matrix computation on big datasets.
    * Götz, M., Debus, C., Coquelin, D., Krajsek, K., Comito, C., Knechtges, P., ... & Streit, A. (2020, December). HeAT–a distributed and GPU-accelerated tensor framework for data analytics. In 2020 IEEE International Conference on Big Data (Big Data) (pp. 276-287). IEEE.

- In the same vein, please correctly credit people for proposing communication schemes like ring-allreduce or similar. In particular:
    * Rabenseifner, R. (2004). Optimization of collective reduction operations. In Computational Science-ICCS 2004: 4th International Conference, Kraków, Poland, June 6-9, 2004, Proceedings, Part I 4 (pp. 1-9). Springer Berlin Heidelberg.