#  Turnip: A “Nondeterministic” GPU Runtime with CPU RAM Offload

Zhimin Ding
Rice University
<EMAIL_ADDRESS>
Jiawen Yao
Rice University
<EMAIL_ADDRESS>
Brianna Barrow
Rice University
<EMAIL_ADDRESS>
Tania Lorido Botran
Roblox
<EMAIL_ADDRESS>
Christopher Jermaine
Rice University
<EMAIL_ADDRESS>
Yuxin Tang
Rice University
<EMAIL_ADDRESS>Jiehui Li
Rice University
<EMAIL_ADDRESS>Xinyu Yao
Rice University
<EMAIL_ADDRESS>Sleem Mahmoud Abdelghafar
Rice University
<EMAIL_ADDRESS>Daniel Bourgeois
Rice University
<EMAIL_ADDRESS>
###### Abstract

An obvious way to alleviate memory difficulties in GPU-based AI computing is
via _CPU offload_ , where data are moved between GPU and CPU RAM. While CPU
offload is useful, it can greatly slow down a computation due to the
relatively transfer rate between CPU RAM and GPU RAM. To address this,
overlapping memory transfer and compute is a necessity, but this
asynchronicity introduces nondeterminacy, and hence it is impossible to know
beforehand what is the best order of operations. We describe Turnip, which is
a system for running AI computations using CPU offload, designed to handle
this nondeterminacy. The key innovation in Turnip is the compilation of the AI
computation into a dependency graph that gives the Turnip runtime freedom to
run operations such as GPU kernel calls in many different orders; at runtime,
Turnip chooses the best order on-the-fly, in response to real-time events. We
find that Turnip outperforms standard, PyTorch-based systems supporting
constrained GPU RAM by a significant margin, and also avoids out-of-memory
errors in a severely memory constrained environment.

##  1 Introduction

Memory management for modern AI computations is difficult. In the LLaMA large
language model (Touvron et al., 2023), for example, the attention computation
on a sequence of length $n$ produces an intermediate result with $128\times
n^{2}$ floating point numbers. Thus, for a long input sequence of 100,000
tokens, the attention computation will result in 1.2 trillion numbers, which
would require 2.4 terabytes to store at half-precision. It is for these
reasons that out-of-memory errors plague AI programmers (FAQ, ).

CPU offload—where data are moved to CPU RAM for storage—can help. As CPU RAM
is much more inexpensive than GPU RAM and it is possible to install many
terabytes of CPU RAM in a GPU server essentially “for free”,111At current
prices, a single state-of-the-art H100 GPU costs the same as approximately
10TB of CPU RAM—more than 100$\times$ the RAM available on a H100 GPU it makes
sense to leverage CPU RAM to temporarily store data. This idea has been
explored in several systems such as pofo (Beaumont et al., 2021), AutoTM
(Hildebrand et al., 2020), SwapAdvisor (Huang et al., 2020), Capuchin (Peng et
al., 2020), and POET (Patil et al., 2022). These systems view a GPU
computation as a dataflow graph, and plan how to fit the computation into GPU
RAM by making use of CPU RAM offload.

While CPU offload is an obvious idea, it can greatly slow down a computation,
due to the relatively slow transfer rate between CPU RAM and GPU RAM. Thus,
any system for CPU offload must ensure that when such a transfer happens, no
computation is blocked waiting for the transfer to finish.

In this paper, we propose Turnip (short for “nondeTerministic gpU RuNtime wIth
cPu offload”) which is a runtime for multi-GPU servers, designed to
systematically support CPU RAM offload. The key innovation of Turnip is its
combination of a pre-computed memory access plan called a memgraph with a
“nondeterministic,” event-driven system runtime. A memgraph is a dependency
graph where vertices represent tasks (such as the execution of a GPU kernel to
perform a small part of of attention computation in a layer of a large
language model) and edges represent data or memory dependencies. Any execution
order that respects the dependencies in the memgraph is valid, and tasks are
dispatched at any time that their dependencies have been met and the
appropriate resources are free. Thus, two executions of the same memgraph may
lead to different sequences of operations being executed on a GPU, or
different sequences of tensors being paged to CPU RAM—hence the non-
determinacy. However, the dependencies in the memgraph are such that the final
output is always _correct_ , no matter the execution order. Turnip’s event-
driven, fully asynchronous runtime is unique. Because operations can be
dispatched whenever the dependencies are fulfilled and are not constrained to
any specific ordering, it lowers the chance that any GPU will be stalled
waiting for a memory transfer to complete. If one task cannot run due to an
un-met dependency in the memgraph, it is possible that there is another task
that _can_ run.

The key technical challenge is how to effectively build a memgraph with as few
dependencies as possible, to allow the runtime as much freedom as possible to
dispatch operations so that it is never blocked, waiting for a memory transfer
to complete. Turnip builds a memgraph by simulating an execution of the
computation, mapping tensors to GPU memory locations and, when necessary
adding edges that represent memory dependencies, as well as offload and reload
operations.

##  2 Why Is Non-Determinacy of Execution Order Crucial?

The design of Turnip is based on a simple hypothesis: When running a GPU-based
computation that utilizes CPU RAM, asynchronous operations such as offload and
reload will have a seemingly nondeterministic running time that is difficult
to pre-plan for. The system runtime must accommodate the resulting non-
determinacy, or else performance can suffer.

Figure 1: A simple memgraph.

Consider Figure 1, which depicts a memgraph for a single GPU system with an
offload (data movement from GPU RAM to CPU RAM) and a reload (from CPU RAM to
GPU RAM). Vertices are operations (for example, GPU kernel calls) that produce
data. Black edges indicate data or consumption dependencies; red edges
indicate memory dependencies (memgraphs will be described in detail in Section
4). Note the memory dependency from $\texttt{offload}_{1}$ to 4. This exists
because the output of 4 will be written to the location of the output of 1,
and so 4 cannot execute until the $\texttt{offload}_{1}$ completes. The data
dependency from $\texttt{reload}_{3}$ to 5 exists because the kernel
associated with vertex 5 will consume the reloaded tensor.

Imagine that a system has executed GPU kernels associated with vertices 1 and
2. It is currently executing the $\texttt{offload}_{1}$ and the
$\texttt{reload}_{3}$. At this point, _it is impossible to know which kernel
should run next_ (4 or 5) as this depends on which memory transfer finishes
first. Ideally, _this decision will be made at runtime_. If the system
deterministically decides to run 4 before 5 at compile time, and the
$\texttt{reload}_{3}$ finishes first, the GPU will sit, idle, waiting for the
$\texttt{offload}_{1}$ to complete. This is why a special-purpose,
“nondeterministic” runtime is needed: to properly handle the non-determinacy
induced via the addition of memory operations.

##  3 Related Work

There are two approaches taken by systems dealing with limited GPU memory.
Some, like Turnip, accept an abstracted version of a generic GPU computation.
Other systems are more specifically targeted to certain categories of models,
optimization algorithms, or to specific tasks such as training or inference.
Unlike Turnip, none of these existing systems consider the effect of the
nondeterminism of offload and reload operations on system performance, nor do
any focus on the system runtime.

Using the first, more general approach, are systems that accept a generic
dataflow graph and, like Turnip, plan for execution in limited memory: pofo
(Beaumont et al., 2021), AutoTM (Hildebrand et al., 2020), SwapAdvisor (Huang
et al., 2020), Checkmate (Jain et al., 2020), Capuchin (Peng et al., 2020),
and POET (Patil et al., 2022) all assume an input dataflow graph for a machine
learning computation, and then plan for execution in limited memory. Checkmate
considers only tensor re-materialization, whereas POET, pofo, and Capuchin
consider re-materialization and offload; AutoTM and SwapAdvisor consider only
offload.

The more targeted approach is taken by the DeepSpeed project (Deepspeed, ) and
the various ZeRO optimizations. For transformers and other, similar models,
DeepSpeed inference (which includes ZeRO-Inference) (Aminabadi et al., 2022)
has two key ideas. First, DeepSpeed inference “offload[s] some activation from
GPU to CPU memory while not in use.” Second, DeepSpeed inference “pins the
model weights either in DRAM (if large enough) or NVMe, and streams each layer
into GPU memory for computation when needed.” FlexGen (Sheng et al., 2023)
seeks to use a variety of methods to speed transformer inference given limited
hardware, including model weight offload to CPU, quantization (Yao et al.,
2022; Frantar et al., 2022), and sparse attention (Child et al., 2019). The
latter two ideas are orthogonal to the ideas in this paper. For CPU offload,
FlexGen optimizes a “zig-zag” block scheduling that works through a tansformer
layers and sequences in the batch, offloading and reloading the KV-cache (Pope
et al., 2023) and model weights. PagedAttention (Kwon et al., 2023) deals with
low memory utilization in transformers, developing a paging system for the KV-
cache.

Figure 2: A decomposition of matrix multiplication.

ZeRO-Offload (Ren et al., 2021) is a comprehensive solution for limited-memory
training that can be seen as primarily using CPU RAM for running the ADAM
optimizer, moving weights to GPU RAM on a carefully-controlled schedule. ZeRO-
Offload is an enhancement on ZeRO (Rajbhandari et al., 2020), which is
designed to be memory-efficient, partitioning both the optimizer and the data
across multiple GPUs. ZeRO-Infinity (Rajbhandari et al., 2021) is similar, and
includes a CPU offload engine, as well as tiling of operators to utilize the
RAM of multiple GPUs.

##  4 taskgraphs and memgraphs in Turnip

Figure 3: Example taskgraph consisting of six GPU kernel calls and three GPU-
to-GPU transfers.

Figure 4: Possible mapping of the output of all of the operations in Figure 4
to memory locations.

Turnip takes is input a taskgraph. A taskgraph is a dataflow graph (a
directed, acyclic graph) that describes how to perform multi-GPU computations.
In a taskgraph, edges represent data flow, and vertices represent operations
over tensors. A vertex without any inputs (called an _input vertex_) is
associated with an input tensor. An operation associated with a non-input
vertex may be either a kernel call that is to be executed on a specific GPU,
or a GPU-to-GPU data transfer.

Turnip is agnostic as to how the taskgraph is created; it could, for example
be created using a framework such as FlexFlow (Jia et al., 2019) or Alpa
(Zheng et al., 2022). Consider a matrix multiplication
$\textbf{X}\times\textbf{Y}$, and assume we wish to execute this matrix
multiplication on three GPUs. To produce a taskgraph, a framework such as
FlexFlow may choose to decompose this matrix multiplication as depicted in
Figure 2, perhaps corresponding to the taskgraph of Figure 4.

Given such a taskgraph, Turnip first compiles the taskgraph into a memgraph,
which it will eventually execute. Like a taskgraph, a memgraph is also a
directed acyclic graph. Every vertex in the original taskgraph will be present
in a corresponding memgraph. Further, the compilation process may add
additional offload and reload operations that move memory from GPU RAM to CPU
RAM, and vice versa. During the compilation process, the output associated
with every vertex in the memgraph is mapped to a memory location. Unlike the
input taskgraph, the memgraph is not a dataflow graph; it is a dependency
graph. If there is an edge from $v_{1}$ to $v_{2}$, it means that $v_{2}$
depends on $v_{1}$ and $v_{2}$ may not execute until after $v_{1}$ has
executed. In a memgraph, there are two types of dependencies. One is a data
dependency, which is inherited from the taskgraph (or is created via the
addition of an offload or reload; see below). The second is a memory
dependency, which is added to ensure that there are no race conditions in the
graph. A _race condition_ occurs when there is some vertex for which two valid
executions of the graph may produce a different output. This can happen when
two vertices write to the same memory location, and it is possible for a third
vertex to read either outupt, depending upon the execution order.

Figure 5: memgraph corresponding to Figure 4.

Let us illustrate a possible compilation of the taskgraph of Figure 4 to a
memgraph. Imagine that our three GPUs each have five memory locations, and for
simplicity, each tensor is the same size and occupies exactly one memory
location. During compilation, the tensor associated with the output of each
operation in the taskgraph is assigned to a memory location, as depicted in
Figure 4. GPU 1 must deal with seven tensors total (two input tensors and five
additional tensors that are created via the execution of some operation), and
we cannot fit all seven of those tensors in memory, given our five locations.
Thus, the tensors output by operations A and 4 are both mapped to GPU1-Loc1,
and the tensors output by operations 1 and 8 are both mapped to GPU1-Loc3.

Figure 6: Possible mapping of tensors to GPU RAM.

A corresponding memgraph is shown in Figure 5. Note that two new edges
representing memory dependencies have been added. These edges guarantee that
the graph is free of race conditions. Specifically, a graph will be free of
race conditions, if, whenever the outputs of vertices $v_{1}$ and $v_{2}$ have
both been mapped to the same memory location, either $v_{1}$ _safely
overwrites the result of_ $v_{2}$, or $v_{2}$ _safely overwrites the result
of_ $v_{1}$. We say that “$v_{1}$ safely over-writes the result of $v_{2}$” if
and only if, for every $v_{3}$ that consumes the output of $v_{2}$, there is a
memory dependency from $v_{3}$ (or some descendent of $v_{3}$) to $v_{1}$ (or
to some ancestor of $v_{1}$). Why? If $v_{1}$ is to safely over-write the
result of $v_{2}$, we need to ensure that $v_{1}$ cannot execute until all of
the consumers of $v_{2}$ have finished execution—such memory dependencies
ensure this.

For example, from Figure 4 we see that the output of vertex 4 is mapped to the
same location as the output of vertex A. In the associated memgraph of Figure
4, to ensure that 4 safely over-writes the result of A, we add a memory
dependency from 3 (the only consumer of A) to 4. From Figure 4 we also see
that the results of 1 and 8 are mapped to the same location. To ensure that 8
safely over-writes the result of 1, we add a memory dependency from 3 (the
only consumer of 1) to 8. Note that this memory dependency is shown as a
dashed line; this indicates that it is _superfluous_ , as there is already a
data dependency from 3 to 8, so this memory dependency is not needed for
correctness.

Figure 7: memgraph with less GPU RAM.

Things can become more intricate if the memory is more constrained. Consider
the case where we have only four memory locations on each GPU, and we wish to
compile the same taskgraph. One possible mapping of the vertices taskgraph of
Figure 4 to memory locations for GPU 1 is shown in Figure 6; the associated
taskgraph is shown in Figure 7. Note in particular the addition of an offload-
reload pair. Both the offload and the reload are new operations that are added
to the memgraph during compilation, to facilitate execution in memory-
constrained scenarios. We can always compile $v_{1}\rightarrow v_{2}$ in a
taskgraph to
$v_{1}\rightarrow\texttt{offload}_{v_{1}}\rightarrow\texttt{reload}_{v_{1}}\rightarrow
v_{2}$ in a memgraph. After the $\texttt{offload}_{v_{1}}$, the result of
$v_{1}$ takes up no GPU memory, but it cannot be used until the
$\texttt{reload}_{v_{1}}$, where it is again mapped to a GPU memory location.
The reason for the inclusion of the offload-reload pair in this case is that
it allows the result of A to be removed from GPU RAM for a time. Thus, vertex
4 can execute and write its result on top of the result of A, which is
subsequently reloaded so that vertex 3 can be executed.

In Figure 6 we see that there are four pairs of vertices whose results are
mapped to the same GPU memory locations, and so memory dependencies must been
added to the memgraph to ensure that there are no race conditions. Consider A
and 4, which are both mapped to GPU1-Loc1. To ensure that 4 safely over-writes
the result of A, we have a memory dependency from the
$\texttt{offload}_{v_{1}}$ (the only consumer of A) to 4. Or, consider
$\texttt{reload}_{v_{1}}$ and 2, which are both mapped to GPU1-Loc4. To ensure
that the $\texttt{reload}_{v_{1}}$ safely over-writes the result of 2, there
is a memory dependency from the only consumer of 2 (vertex 4) to the
$\texttt{reload}_{v_{1}}$.

##  5 The Turnip Execution Engine

Once a memgraph has been produced, it is executed by the Turnip engine using a
nondeterministic, event-based framework. As soon as a GPU is unused or a
tensor is ready to be offloaded to RAM, the Turnip runtime can immediately
assign any available work to the GPU or begin the transfer, without regard to
the overall state of the computation. Also note that there are no calls to
memory-management routines such as cudaMalloc or cudaFree during memgraph
execution, as memory management is no longer dynamic. Tensor placement is pre-
determined before execution, and if dependencies are respected, there can be
no memory corruption due to race conditions.

To execute the memgraph, Turnip runs a central event processing loop, that
repeatedly processes callback functions that are called response to completion
of the work associated with a memgraph vertex (completion of a GPU-to-GPU
transfer, completion of the GPU kernel, or completion of an offload or
reload). When a vertex completes and a callback is invoked, the event loop
checks to see if any other vertex can be executed. That is, it searches for a
vertex $v_{1}$ where (a) all vertices $v_{2}$ with an edge $v_{2}\rightarrow
v_{1}$ in the memgraph have also completed; (b) if $v_{1}$ is a kernel call,
then the GPU $v_{1}$ is assigned is currently free. When the event loop finds
such a vertex, it launches it, and searches for another such vertex. When it
can find no such executable vertex, it goes to sleep until woken by another
callback.

##  6 Building a memgraph

The key technical question we address in this paper is: How to construct a
memgraph from a taskgraph? The primary requirement for the compilation process
is _correctness_. Correctness requires that (a) every data dependency present
in the taskgraph is also present in the memgraph, _or_ is replaced with a
sequence of offload-reload operations;222So, for example, if $v_{1}\rightarrow
v_{2}$ is present in the taskgraph, we may have
$v_{1}\rightarrow\texttt{offload}_{v_{1}}\rightarrow\texttt{reload}_{v_{1}}\rightarrow
v_{2}$ in the memgraph (b) there are no race conditions in the memgraph; (c)
the memgraph has no cycles. In addition, it is desirable for the memgraph to
be performant. A memgraph will not be performant if memory dependencies
severely constrain the execution order of vertices. Such constraints may
reduce parallelism and GPU utilization.

BuildMemgraph: Inputs: taskgraph, sorted list of taskgraph vertices $V=\langle
v_{1},v_{2},...,v_{n}\rangle$; Outputs: memgraph, GPU memory location
$v_{i}.\texttt{loc}$ for $i\in\\{1...n\\}$

———————————————————————-

$\texttt{Evicted}\leftarrow\\{\\}$; $\texttt{execHzn}\leftarrow 1$;
$\texttt{allocHzn}\leftarrow 1$;

while execHzn $\leq n$ do

if allocHzn $<=n$ and
$(v_{\texttt{allocHzn}}.\texttt{loc}\leftarrow\texttt{simMalloc}(v_{\texttt{allocHzn}}))\neq-1$
then

/*successfully allocated space for future result*/

allocHzn += $1$

else if $\texttt{allocHzn}=\texttt{execHzn}$ then

/* unable to allocate for next execution w/o evict*/

$v_{\texttt{allocHzn}}.\texttt{loc}\leftarrow\texttt{simMallocOffld}(v_{\texttt{allocHzn}})$

allocHzn += $1$

else

/* simulate execution of the next vertex */

/* first, compute set of vertices exec depends on */

$\texttt{Deps}\leftarrow\\{v\textrm{ s.t. edge }v\rightarrow
v_{\texttt{execHzn}}\in\textsc{{taskgraph}}{}\\}$

for $v\in\texttt{Deps}$ do

/* reload dependency if evicted */

if $v\in\texttt{Evicted}$ then

$v.\texttt{loc}\leftarrow\texttt{simMallocForceReld}(v)$

end if

/* if dependency won’t be used again, free it */

if not $\exists(\texttt{fut}>\texttt{execHzn}\texttt{ s.t. edge }v\rightarrow
v_{\texttt{fut}}\in\textsc{{taskgraph}}{})$ then

$\texttt{simFree}(v)$

end if

add edge $v\rightarrow v_{\texttt{execHzn}}$ to memgraph

end for

execHzn += $1$

end if

end while

Figure 8: Building a memgraph via execution simulation.

Our basic tactic during compilation is to rely on a simulated execution of the
taskgraph to generate the memgraph. Given a serial ordering of the vertices in
the taskgraph that respects all dependencies (so that if $v_{1}\rightarrow
v_{2}$ is in the taskgraph, $v_{1}$ is before $v_{2}$ in the ordering) we
simulate its execution, making calls to special variants of malloc and free
that do not actually allocate GPU RAM, but instead maintain a map of used and
free RAM slots on the GPU that is the target of the compilation. These
implementations also maintain a history of which tensors occupied which
positions in simulated GPU RAM, to correctly generate memory dependencies. As
the simulation runs, the memgraph is constructed. Calls to the special malloc
implementations associate memgraph vertex outputs to GPU memory locations
(effectively producing the mappings depicted in Figure 4 and Figure 6).
Whenever a call to the malloc variant fails because there is not enough GPU
RAM, an offload vertex must be added to the memgraph. Whenever it is time to
simulate execution of a taskgraph vertex but one of the inputs is not in the
simulated GPU RAM, then a memory location for the corresponding reload vertex
is allocated, and a data dependency on that reload is added to the memgraph.

As the simulation runs, there are two _horizons_ , or counters that mark
progress through the serialized taskgraph. The first is the allocHzn. Every
vertex in the taskgraph that is older than the allocHzn has had a space
allocated for it. The second is the execHzn. Every vertex in the taskgraph
that is older than the execHzn has been “run” according to the simulation. To
ensure a high-quality memgraph, our compilation algorithm greedily tries to
push the allocHzn as far as possible past the execHzn. Intuitively, this will
produce fewer constraints in the resulting memgraph. A kernel associated with
a vertex cannot run until it has GPU RAM to write its output. If this GPU RAM
is available very early in the simulation, then it gives the Turnip event
processing loop more freedom to choose a vertex execution order that does not
exactly match the simulated ordering, generating more opportunities to run
available kernels while waiting for memory transfers.

The overall algorithm, BuildMemgraph, is given above in Figure 8. Note that
this variant of the algorithm assumes each tensor takes up exactly one slot in
GPU RAM. In the ”real life” case where tensors are variably-sized, the
algorithm does not change appreciably—specifically, in the variably-sized
case, freeing space for a tensor can evict a variable number of tensors to CPU
RAM—but assuming uniformly-sized tensors simplifies the presentation.

At the highest level, the algorithm operates by first checking to see if it
can allocate space for the vertex at the current allocation horizon,
$v_{\texttt{allocHzn}}$. If it cannot, the algorithm makes sure there is space
available for the output of the next vertex to be executed (the only way there
is not space is if $\texttt{allocHzn}=\texttt{execHzn}$ and the last
allocation failed; this implies it is time to execute $v_{\texttt{execHzn}}$
and we just failed to allocate space for it). If there is space, the
simulation “executes” $v_{\texttt{execHzn}}$.

simMalloc: Input: vertex $v$; Output: GPU memory slot for $v$

——————————————————

find open slot for $v$; return -1 if none

return slot if no previous occupant

$v^{\prime}\leftarrow$ last owner of slot for $v$

$\texttt{Deps}\leftarrow\\{v^{\prime\prime}\textrm{ s.t. edge
}v^{\prime}\rightarrow v^{\prime\prime}\in\textsc{{taskgraph}}{}\\}$

for $v\in\texttt{Deps}$ do

add edge $v^{\prime\prime}\rightarrow v$ to memgraph

end for

return slot

simMallocForceReld: Input: vertex $v$; Output: GPU memory slot for $v$

——————————————————

remove $v$ from Evicted

$\texttt{slot}\leftarrow\texttt{simMalloc}(v)$

if $\texttt{slot}\neq-1$ then

return slot

end if

return $\texttt{simMallocOffld}(v)$

simMallocOffld: Input: vertex $v$; Output: GPU memory slot for $v$

——————————————————

find GPU RAM slot for $v$ and determine victim (current occupant of slot)
$v^{\prime}$

add sequence
$v^{\prime}\rightarrow\texttt{offload}_{v^{\prime}}\rightarrow\texttt{reload}_{v^{\prime}}$
to memgraph

add edge $\texttt{offload}_{v^{\prime}}\rightarrow v$ to memgraph

$\texttt{Deps}\leftarrow\\{v^{\prime\prime}\textrm{ s.t. edge
}v^{\prime}\rightarrow v^{\prime\prime}\in\textsc{{taskgraph}}{}$ and
$v^{\prime\prime}$ comes before $v$ in $V\\}$

for $v\in\texttt{Deps}$ do

add edge $v^{\prime\prime}\rightarrow v$ to memgraph

end for

rename all instances of $v^{\prime}$ in taskgraph to
$\texttt{reload}_{v^{\prime}}$

add $\texttt{reload}_{v^{\prime}}$ to Evicted

return slot

Figure 9: simMalloc variants used in memgraph construction.

There are four memory management subroutines used by the algorithm: three
variants on malloc (simMalloc, simMallocForceReld, and simMallocOffld) and one
variant on free called simFree. Like a traditional malloc, simMalloc finds an
open slot for the allocation, but it also adds the memory dependencies to the
memgraph necessary to ensure that the vertex $v$ that will occupy the slot
will safely overwrite the previous occupant of the slot. simMallocForceReld is
like simMalloc, but it is used in the case when a vertex must be reloaded
because it is going to be used immediately, and hence the allocation cannot
fail. simMallocOffld is a variant of simMalloc that cannot fail, as it finds a
victim to offload to ensure the success of the allocation for vertex $v$,
adding the offload-reload sequence to the memgraph. Crucially, it renames all
instances of the victim $v^{\prime}$ in the taskgraph to refer to
$\texttt{reload}_{v^{\prime}}$. In this way, all “future” accesses to
$v^{\prime}$ will refer, in fact, to its reloaded version. The routine also
adds a memory dependency from the $\texttt{offload}_{v^{\prime}}$ to $v$, as
we cannot execute $v$ until the $\texttt{offload}_{v^{\prime}}$ has taken
place, and freeing GPU RAM for use.

##  7 Correctness of memgraph Construction

By construction, all edges present in the taskgraph are present in the
memgraph, as when a vertex is “executed” during the simulation, all of its
incoming data dependencies are added to the memgraph.

To ensure correctness with respect to race conditions and the absence of
cycles, the algorithm relies on the total ordering of vertices in $V$, the
list input into BuildMemgraph. Further, all new offload and reload vertices
produced by the compilation process also have a consistent placement in this
ordering. Consider the case when simMallocOffld or simMallocForceReld are
called to obtain data necessary to execute vertex $v$, and an offload or a
reload is produced. All of those offloads and reloads take place just before
$v$, with all of the offloads happening just before the reloads.

This ordering ensures that there can be no race conditions in the output
memgraph: if the outputs of $v_{1}$ and $v_{2}$ are mapped to the same memory
location and $v_{1}$ comes before $v_{2}$ in the ordering, then the
BuildMemgraph algorithm ensures that $v_{2}$ safely overwrites the result of
$v_{1}$. Consider the implementation of simMalloc. Whenever a tensor produced
by $v$ is mapped to a memory slot previously occupied by the result of
$v^{\prime}$, we add edges to ensure that every consumer of $v^{\prime}$
executes before $v$. Also, consider simMallocOffld, where $v^{\prime}$ is
offloaded to CPU RAM to accommodate $v$. Here, memory dependencies are added
from the $\texttt{offload}_{v^{\prime}}$ to $v$ and from all consumers of
$v^{\prime}$ to $v$ (when those consumers appeared before $v$ in the list
$V$). Note that the taskgraph is modified so that all “future” consumers of
$v^{\prime}$ will consume $\texttt{reload}_{v^{\prime}}$ rather than
$v^{\prime}$, so they cannot induce a race condition over $v$ and
$v^{\prime}$.

Also, consider why there can be no cycles in the output memgraph: all edges
added to the memgraph point forward in the total ordering. Consider the edges
added by simMalloc. $v$ can only be mapped to the location used by
$v^{\prime}$ if $v^{\prime}$ has been previously free’ed. This implies that
any vertices using the output of $v^{\prime}$ have already been “executed”,
and so come before $v$ in the total ordering. Thus, any edge from a consumer
of $v^{\prime}$ to $v$ must point forward in $V$. Also consider the edges
added by simMallocOffld. A similar argument holds here, as we explicitly only
add edges that point forward in $V$.

##  8 Experiments

Figure 10: Time taken for LLaMA first-token (prefill) inference. “OOM” is out-
of-memory. Figure 11: Training time per batch for LoRA training, 7B param
model, 4$\times$P100 GPU server.

Our experiments evaluate the ability of Turnip to deal with Meta AI’s LLaMA
large language model (LLM) (Touvron et al., 2023), with severely constrained
memory. LLM training and inference are chosen as representative, challenging
computational workloads encountered by modern ML systems, particularly
difficult given the large memory footprint. We assessed Turnip’s performance
on 7 billion and 65 billion parameter models.

Turnip is implemented in C++, with most GPU kernels generated using Nvidia
cuTensor. Experiments were conducted on two machines: (i) an older 4 $\times$
P100 GPU server (16 GB RAM each) and 22, 64GB DDR4 2666MHz CPU RAM modules,
for a total of 1.3TB of RAM, and (ii) an Amazon Web Services p4d.24xlarge
instance, equipped with eight A100 GPUs (40 GB RAM each) and 1.15TB of RAM. We
were particularly interested in seeing the ability of Turnip to operate in a
difficult environment with extremely limited GPU RAM, hence the P100 GPUs,
with only 64 GB of GPU RAM total on the server. A key question is: can
software help bridge the gap—particularly the lack of GPU RAM—between older
and newer hardware? The tasks evaluated are:

(1) LLaMA first token inference. Our experiments target “first token”
inference (also known as ”prefill”): How long does it take to produce the
first output token, given an input prompt? Prefill is exceedingly expensive in
terms of the memory required, scaling quadratically with the size of the
prompt. On both machines, we run Turnip, ZeRO Inference (Aminabadi et al.,
2022) (using weight partitioning and model weight offload), and FlexGen (Sheng
et al., 2023). Note that these are PyTorch-based systems, whereas Turnip is
not. For FlexGen, we use two settings: both use 50/50 CPU/GPU RAM for the
model weight cache; one uses full CPU offload for activations (“Act. off.”),
and one does not. As an ablation, we also run Turnip (“Fixed execution”),
constrained so that it executes vertices in the order of the simulation used
to build the memgraph (thus nondeterminacy is mostly removed from the runtime;
see the discussion below). For each option, we test the following prefill
input sequence lengths: 1K, 2K, 4K, 8K, 16K, and 32K tokens. All model weights
and computations were performed at half precision (though FlexGen uses even
lower precision for some computations). A subset of the results are given in
Figure 10 (the rest are in the Appendix). For each experiment, we give the
time required to produce the first token. “OOM” inidcates running out-of-
memory.

(2) LoRA training for LLaMA. We also experiment with LoRA training (Hu et al.,
2021). We use a LoRA rank of 16, and train LoRA adaptors for the $K$, $V$,
$Q$, and feedforward mapping matrices. Here we run Turnip and ZeRO Infinity
(Rajbhandari et al., 2021); ZeRO is executed using all three “stages”
(gradient partitioning, model weight partitioning, optimizer state
partitioning) as well as CPU offload. Both Turnip and ZeRO use checkpointing
during the forward pass to reduce the memory footprint. We again run the
Turnip ablation, where execution order is constrained. We measure the time it
takes to run the forward and backward pass for one batch, with varying batch
sizes (1, 4, 16, 64) and sequence lengths (1K, 2K, 4K, and 8K). All model
weights are stored as single precision (32 bits). A subset of the results are
given in Figure 11 (other results are in the appendix).

Discussion. Turnip performs significantly better than ZeRO Inference, FlexGen,
and ZeRO Infinity, both in terms of runtime, and the ability to handle the
severely constrained GPU memory. Among different workloads, Turnip is always
the fastest, and avoids most OOM errors. In the most memory-constrained
inference experiments (P100 GPUs, 65B parameter model) Turnip was the only
method able to go beyond a 2K sequence length without an OOM error. Further,
Turnip was able to train the 65B parameter model in many cases on the P100
server, with only 64GB of GPU RAM (see Figure 13 in the Appendix); ZeRO
Infinity could handle only a single 1K sequence at a time. One key reason for
the robustness is the way in which the sub-computations making up the same,
larger operation are decoupled and executed in Turnip. If multiple small
matrix multiplies are used to implement a single large matrix multiply (as in
Figure 3) they need not run together, en-masse. In a severely memory-
constrained environment, Turnip can force them to be run in sequence and
offloaded, via the introduction of memory dependencies. The other systems do
not do this, and perform layerwise execution of the transformer. This is why
FlexGen with CPU offload does not help. As we are performing inference on a
single sequence, offloading an activation is not useful (FlexGen’s offload
targets CPU offload for the different sequences in a batch, and cannot
partially offload the activations for a single sequence).

A secondary finding is that forcing a particular execution order in Turnip
(“Fixed execution”) resulted in up to a 3$\times$ slowdown compared to the
nondeterministic order. Note that even when nondeterminstic execution is
“turned off,” Turnip is still much faster than the PyTorch-based systems, even
with ample GPU RAM. One reason for this may be that due to the architecture of
Turnip, it is very difficult to completely “turn off” nondeterminism during
aggregation—Turnip does not use a synchronous all-reduce, for example, as the
other systems do, and instead uses a series of asynchronous, partial
summations. This remaining nondeterminism may help explain the still-
significant performance advantage of the “Fixed execution” option. Additional
experimental results and discussion are given in the Appendix.

##  9 Conclusions

We present Turnip, a system and runtime for executing memory-intensive
computations on GPU servers, supporting CPU offload. The key innovation of
Turnip is its reliance on a “nondeterministic” runtime where a dependency
graph is used to guide execution of the computation; any order of operations
is acceptable, if the dependencies are respected. We argue that this is
necessary when CPU RAM offload is used, or else the system will often be
stalled, waiting for CPU-to-GPU transfers.

Limitations. As currently implemented, the biggest limitation of Turnip is
that the input computation (in the form of a taskgraph) must be static, and
known completely beforehand, so that the memgraph can be constructed. This is
not of consequence during model training, but can be an issue during any
recursive, generative AI computation. This includes LLM inference, where the
next token must repeatedly be generated and the $KV$-cache increases in size.
There are some naive solutions to this (such as pre-compiling a memgraph for a
specific number of token generations in the case of an LLM) but more work will
be necessary to create a satisfactory solution to recursive computation.

Another limitation is that while our experiments did show that Turnip has
certain performance advantages, it is not possible to be absolutely sure where
those advantages come from. Our ablation shows that a fixed execution order
slows down Turnip, suggesting that nondeterminism is important. But unlike
ZeRO and FlexGen, Turnip was implemented from the ground up and does not rely
on Python or PyTorch—that fact alone could account for some of the performance
differences.

## Acknowledgments and Disclosure of Funding

This work was supported by an NIH CTSA award No. UL1TR003167, by the NSF under
grant numbers 1918651, 2008240, 2131294, and 2212557, by the US DOT Tier-1 UTC
CYBER-CARE grant #69A3552348332, and by a gift from the Roblox Corporation.

## References

  * Aminabadi et al. (2022) Reza Yazdani Aminabadi, Samyam Rajbhandari, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Olatunji Ruwase, Shaden Smith, Minjia Zhang, Jeff Rasley, et al.  Deepspeed-inference: enabling efficient inference of transformer models at unprecedented scale.  In _SC22: International Conference for High Performance Computing, Networking, Storage and Analysis_ , pp. 1–15. IEEE, 2022.
  * Beaumont et al. (2021) Olivier Beaumont, Lionel Eyraud-Dubois, and Alena Shilova.  Efficient combination of rematerialization and offloading for training dnns.  _Advances in Neural Information Processing Systems_ , 34:23844–23857, 2021.
  * Child et al. (2019) Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever.  Generating long sequences with sparse transformers.  _arXiv preprint arXiv:1904.10509_ , 2019.
  * (4) Deepspeed.  Deepspeed.  https://www.microsoft.com/en-us/research/project/deepspeed/.  Accessed: 2024-05-12.
  * (5) Pytorch FAQ.  Pytorch frequently asked questions.  https://pytorch.org/docs/stable/notes/faq.html.  Accessed: 2024-05-12.
  * Frantar et al. (2022) Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh.  Gptq: Accurate post-training quantization for generative pre-trained transformers.  _arXiv preprint arXiv:2210.17323_ , 2022.
  * Hildebrand et al. (2020) Mark Hildebrand, Jawad Khan, Sanjeev Trika, Jason Lowe-Power, and Venkatesh Akella.  Autotm: Automatic tensor movement in heterogeneous memory systems using integer linear programming.  In _Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems_ , pp. 875–890, 2020.
  * Hu et al. (2021) J. Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, and Weizhu Chen.  Lora: Low-rank adaptation of large language models.  _ArXiv_ , abs/2106.09685, 2021.  URL https://api.semanticscholar.org/CorpusID:235458009.
  * Huang et al. (2020) Chien-Chin Huang, Gu Jin, and Jinyang Li.  Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping.  In _Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems_ , pp. 1341–1355, 2020.
  * Jain et al. (2020) Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Joseph Gonzalez, Kurt Keutzer, and Ion Stoica.  Checkmate: Breaking the memory wall with optimal tensor rematerialization.  _Proceedings of Machine Learning and Systems_ , 2:497–511, 2020.
  * Jia et al. (2019) Zhihao Jia, Matei Zaharia, and Alex Aiken.  Beyond data and model parallelism for deep neural networks.  _Proceedings of Machine Learning and Systems_ , 1:1–13, 2019.
  * Kwon et al. (2023) Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica.  Efficient memory management for large language model serving with pagedattention.  In _Proceedings of the 29th Symposium on Operating Systems Principles_ , pp. 611–626, 2023.
  * Patil et al. (2022) Shishir G Patil, Paras Jain, Prabal Dutta, Ion Stoica, and Joseph Gonzalez.  Poet: Training neural networks on tiny devices with integrated rematerialization and paging.  In _International Conference on Machine Learning_ , pp. 17573–17583. PMLR, 2022.
  * Peng et al. (2020) Xuan Peng, Xuanhua Shi, Hulin Dai, Hai Jin, Weiliang Ma, Qian Xiong, Fan Yang, and Xuehai Qian.  Capuchin: Tensor-based gpu memory management for deep learning.  In _Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems_ , pp. 891–905, 2020.
  * Pope et al. (2023) Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean.  Efficiently scaling transformer inference.  _Proceedings of Machine Learning and Systems_ , 5, 2023.
  * Rajbhandari et al. (2020) Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He.  Zero: Memory optimizations toward training trillion parameter models.  In _SC20: International Conference for High Performance Computing, Networking, Storage and Analysis_ , pp. 1–16. IEEE, 2020.
  * Rajbhandari et al. (2021) Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He.  Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning.  In _Proceedings of the international conference for high performance computing, networking, storage and analysis_ , pp. 1–14, 2021.
  * Ren et al. (2021) Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He.  $\\{$Zero-offload$\\}$: Democratizing $\\{$billion-scale$\\}$ model training.  In _2021 USENIX Annual Technical Conference (USENIX ATC 21)_ , pp. 551–564, 2021.
  * Sheng et al. (2023) Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, Ion Stoica, and Ce Zhang.  Flexgen: High-throughput generative inference of large language models with a single gpu.  In _International Conference on Machine Learning_ , pp. 31094–31116. PMLR, 2023.
  * Touvron et al. (2023) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_ , 2023.
  * Yao et al. (2022) Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He.  Zeroquant: Efficient and affordable post-training quantization for large-scale transformers.  _Advances in Neural Information Processing Systems_ , 35:27168–27183, 2022.
  * Zheng et al. (2022) Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P Xing, et al.  Alpa: Automating inter-and $\\{$Intra-Operator$\\}$ parallelism for distributed deep learning.  In _16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)_ , pp. 559–578, 2022.

##  Appendix A Additional Experimental Results and Discussion

Figure 12: Time taken for LLaMA first-token (prefill) inference. 7B
parameters, 4 $\times$ P100. “OOM” is out-of-memory. Figure 13: Training time
per batch for LoRA training, 65B param model, 4 $\times$ P100 GPU server.
Figure 14: Training time per batch for LoRA training, 7B param model, 8
$\times$ A100 GPU server.  Figure 15: Training time per batch for LoRA
training, 65B param model, 8 $\times$ A100 GPU server.

Here we give some additional experimental results not found in the main body
of the paper, as well as some additional discussion. Figure 12 gives inference
results for the P100 server (7B params). Figure 13 gives training results for
the 65B parameter model, on the P100 server. Figure 14 gives results for
training the 7B parameter model on the A100 server. Figure 15 gives results
for training the 65B parameter model on the A100 server.

On training, we were initially surprised to find that there is a relatively
small difference in performance between the 4 $\times$ P100 and the 8 $\times$
A100 servers. For example, for batch size 64, 1K sequence length, 7B parameter
model, the P100 server took 174 seconds per batch (for Turnip) and 35 seconds
per batch on the A100 server. That is, the P100 server is only around
$5\times$ slower, despite the fact that it has half the number of GPUs, and
the GPUs are two generations older. However, upon reflection, this is not
_that_ surprising. Much of the increase in compute power in the never GPUs is
related to increased 16-bit performance, which is not particularly useful for
training. One immediate effect of this is that it magnifies the importance of
software, as the performance of the servers for training is not _that_
different—the real problem is the lack of GPU RAM on the P100 server—and a
software that can deal with this lack of GPU RAM and cut out as much
inefficiency as possible can bridge the gap between the two hardwares. For
example, note that Turnip processes a batch of 16, 1K sequences on the P100
server (7B parameter model) in 43 seconds. ZeRO Infinity processes the same
16, 1K sequences on the A100 server in 48 seconds—ZeRO Infinity is actually
five seconds slower per batch on the newer, A100 server, compared to Turnip on
the P100 server.

Note that on the inference experiments, FlexGen has a relatively flat curve as
the sequence length increases. It starts out slower than ZeRO at small
sequence lengths, but does not slow down as much at longer sequence lengths.
One reason for this may be the various methods embedded in FlexGen (such as a
faster attention computation) for speeding up the computation. At smaller
sequence lengths, the overhead is high, but the overhead may become less of a
factor at larger sequences. Still, Turnip is always faster than both ZeRO and
FlexGen—though to be fair, FlexGen specifically targets batch processing, and
all of our experiments are single-sequence. Also note that for all
experiments, Turnip implements a “vanilla” computation—we use the LLaMA
attention mechanism as implemented by LLaMA, and use 32-bit computation
everywhere for training, and 16-bit computation everywhere for inference. No
techniques such as quantization are used to save space and/or speed inference.

##  Appendix B Turnip Engine Details

Our execution engine consists of a central event loop that “launches” each
vertex in the memgraph. A vertex can be launched when (1) all dependencies
have been completed and (2) the required resources are obtained. When a vertex
is launched, the corresponding operation is executed and then a provided
callback is called to notify the event loop that the vertex has completed. In
turn, the event loop frees up the obtained resources and keeps track of when
vertices complete execution so that subsequent vertices can be launched. In
practice, when launched, a vertex will execute one or more asynchronous CUDA
operations on CUDA stream and will then call cudaStreamAddCallback. As such,
every vertex requires as a resource a stream, where a single stream can only
be used by a single launched vertex at a time. We use 5 streams per GPU.
offload, Reload and inter-GPU communication vertices will call
cudaMemcpyAsync. For CPU storage, we allocate a single, large contiguous block
of memory with cudaHostAlloc with flags cudaHostAllocPortable and
cudaHostAllocWriteCombined. When executing Offload vertices, we allocate into
the CPU storage memory using our custom allocator; when executing Reload
vertices, we free from our custom allocator. All compute vertices are executed
using either cuTensor functions or hand-written CUDA kernels. An example where
a hand-written CUDA kernel is beneficial is for executing portions of softmax
so that less workspace memory and fewer vertices would be required. Two
additional resources may be required for computing vertices: workspace memory
as required for executing multiple cuTensor functions and locks around write-
protected memory. As an example, we would execute a summation of $n$ tensors
with $n$ calls to tensor increment sum-into kernels. However, the output
memory would be protected by a resource so that only one sum-into can happen
at a time. This implementation is designed to support non-determinism. We use
CUDA version 11.8.0 and cuTensor version 2.0.1. All other code is C++.

##  Appendix C Ensuring the “Unobtrusiveness” of Dependencies

Even under nondeterministic execution, memory dependencies can cause periods
of time where a GPU is not used, as the GPU is blocked on a memory transfer or
on a dependency related to two tensors being mapped to the same location in
GPU RAM. The likelihood that dependencies added to the memgraph will cause
such stalls can be reduced by a careful implementation of the various malloc
variants. For example, when simMallocOffld searches for a victim, we search
for the victim whose next use is furthest in the future—that is, closest to
the end of $V$ (in the general case where tensors have different sizes and
there may be more than one victim, we seek to maximize the minimum age of any
evicted tensor). When simMalloc finds an open slot for a tensor, it should
choose the available slot whose last use was furthest in the past (closest to
the beginning of $V$).


11institutetext:  University of Science, VNU-HCM, Vietnam 22institutetext:
Vietnam National University, Ho Chi Minh City, Vietnam

# NearbyPatchCL: Leveraging Nearby Patches for Self-Supervised Patch-Level
Multi-Class Classification in Whole-Slide Images

Gia-Bao Le Both authors contributed equally1122 0009-0002-2815-6798 Van-Tien
Nguyen⋆ 1122 0009-0006-3983-0281 Trung-Nghia Le Coresponding author1122
0000-0002-7363-2610 Minh-Triet Tran 1122 0000-0003-3046-3041

###### Abstract

Whole-slide image (WSI) analysis plays a crucial role in cancer diagnosis and
treatment. In addressing the demands of this critical task, self-supervised
learning (SSL) methods have emerged as a valuable resource, leveraging their
efficiency in circumventing the need for a large number of annotations, which
can be both costly and time-consuming to deploy supervised methods.
Nevertheless, patch-wise representation may exhibit instability in
performance, primarily due to class imbalances stemming from patch selection
within WSIs. In this paper, we introduce Nearby Patch Contrastive Learning
(NearbyPatchCL), a novel self-supervised learning method that leverages nearby
patches as positive samples and a decoupled contrastive loss for robust
representation learning. Our method demonstrates a tangible enhancement in
performance for downstream tasks involving patch-level multi-class
classification. Additionally, we curate a new dataset derived from WSIs
sourced from the Canine Cutaneous Cancer Histology, thus establishing a
benchmark for the rigorous evaluation of patch-level multi-class
classification methodologies. Intensive experiments show that our method
significantly outperforms the supervised baseline and state-of-the-art SSL
methods with top-1 classification accuracy of 87.56%. Our method also achieves
comparable results while utilizing a mere 1% of labeled data, a stark contrast
to the 100% labeled data requirement of other approaches. Source code:
https://github.com/nvtien457/NearbyPatchCL

###### Keywords:

Self-supervised learning Contrastive learning Whole-slide image Representation
learning.

##  1 Introduction

Figure 1: In contrastive learning [8], our NearbyPatchCL, and SupCon [16],
positive samples are handled differently. Contrastive learning involves
augmented pairs from the same image. In SupCon, it includes all images with
the same label, while NearbyPatchCL defines positive samples that encompass
all views of the center patch (blue border) and nearby patches (green border).

Histology image analysis is crucial for understanding biological tissues. As
Whole-Slide Images (WSI) become more common, and as cost-effective storage and
fast data transfer networks become available, the creation of large databases
of digitized tissue sections in hospitals and clinics is on the rise.
Currently, the computational pathology community is actively digitizing tissue
slides into WSI for automated analysis. There’s a growing focus on developing
precise algorithms for clinical use, with recent advancements in deep learning
making automatic analysis of WSIs, whether supervised or weakly supervised,
more popular [4, 17, 24, 28, 32].

While there has been notable progress in the automated processing and clinical
use of WSIs, challenges persist due to their gigapixel size. These challenges
often require the use of tile-level processing and multiple instance learning
for predicting clinical endpoints [28, 32, 22]. Additionally, the large size
of WSIs makes the annotation process by human experts cumbersome, requiring
annotated data for algorithm development. The shift to deep learning further
emphasizes the importance of annotations, but some methods explore the use of
pre-trained representations, typically from ImageNet, as an alternative to
generating WSI-specific representations [11, 18, 20].

Self-supervised learning (SSL) methods [8, 1, 10] are gaining popularity for
their ability to acquire competitive, versatile features compared to
supervised methods. SSL involves two steps: unsupervised pre-training on
unlabeled data and supervised fine-tuning on downstream tasks with limited
labeled data. These methods not only require a small amount of labeled data
but also enhance model performance across various histology pathology tasks
[10, 23, 25, 29, 21, 15, 26, 18, 20]. Commonly, these methods break WSIs into
smaller patches, feed them to an encoder, and extract representation features
for downstream tasks. However, imbalanced category annotations within WSIs can
lead to unstable performance, especially when random cropping generates image
patches. While contrastive learning methods [10, 18, 20, 23, 21] have been
applied to address this issue, it is important to note that even contrastive
learning methods are not immune to imbalanced datasets [14]. Consequently,
several SSL methods require further improvements to seamlessly integrate into
clinical practice.

In this paper, we propose a simple yet efficient self-supervised learning
method called Nearby Patch Contrastive Learning (NearbyPatchCL) that treats
adjacent patches as positive samples in a supervised contrastive framework,
which makes the training process more robust (See Fig. 1). To grapple with the
intricate challenge of imbalanced data, we adopt the decoupled contrastive
learning (DCL) loss [30]. The amalgamation of these approaches not only
enhances the overall effectiveness and stability of our methodology but also
engenders a strong representation that holds substantial promise for clinical
integration. Remarkably, our approach yields commendable performance even when
furnished with a scant fraction of labeled data, demonstrating its potential
utility in real-world applications within the medical domain.

For patch-level multi-class evaluation, we adopt and process WSIs from the
public CAnine CuTaneous Cancer Histology (CATCH) [27], resulting in a new
benchmark dataset, namely P-CATCH. Intensive experiments on the newly
constructed P-CATCH dataset demonstrate the superiority of NearbyPatchCL. Our
method achieves the top-1 classification accuracy of 87.56% and significantly
outperforms the supervised baseline and existing SSL methods. The proposed
method also achieves compatible results when using only 1% labeled data
compared to others using 100% labeled data. The source code will be released
upon the paper’s acceptance. Our contributions are summarized in the
following:

  * •

We propose a novel self-supervised method, namely NearbyPatchCL, which
incorporates the modified supervised contrastive loss function by leveraging
nearby patches as positive samples combined with the decoupled contrastive
learning loss for better representation.

  * •

We introduce a new dataset for benchmarking WSI patch-level multi-class
classification methods.

  * •

We perform a comprehensive comparison with state-of-the-art SSL methods to
demonstrate the superior performance of the proposed method.

##  2 Related works

###  2.1 Self-Supervised Learning

SSL methods can be separated into four types based on their learning
techniques. Contrastive learning algorithms (e.g., SimCLR [8] and MoCo [13])
aim to distinguish individual training data instances from others, creating
similar representations for positive pairs and distinct representations for
negative pairs. However, these methods require diverse negative pairs, often
mitigated by using large batches or memory banks. Yeh et al.[30] introduced
decoupled contrastive learning loss, enhancing learning efficiency by removing
the positive term from the denominator. This method can achieve competitive
performance with reduced sensitivity to sub-optimal hyperparameters, without
the need for large batches or extended training epochs. Additionally,
asymmetric networks, such as BYOL [12] and SimSiam [9], exhibit parallels with
contrastive learning methods as they both learn representations of images from
various augmented viewpoints. BYOL uses 2 distinct networks to create such
representations, while SimSiam utilizes Siamese Networks[3]. A distinguishing
factor compared to contrastive methods is that these strategies operate
independently of negative pair incorporation, enabling them to function
effectively even when dealing with small batch sizes. On the other hand,
clustering-based methods, including DeepCluster [5] and SwAV [6], pursue the
discovery of meaningful and compact representations by leveraging the notion
of similarity and dissimilarity between data points. These methods operate
under the assumption that similar data points should be closer together in the
embedding space, while dissimilar points should be far apart. Meanwhile,
feature decorrelation methods, exemplified by VICReg [2] and BarlowTwins [31],
address redundancy among different dimensions of learned features, preventing
collapse or over-reliance on specific dimensions. By reducing redundancy,
these methods enable more reliable and comprehensive representations,
contributing to advancements in self-supervised learning techniques and
achieving results on par with state-of-the-arts on several downstream tasks.

###  2.2 SSL in Digital Pathology Images Analysis

By leveraging the inherent structure and relationships within the data, SSL
techniques can learn rich and meaningful representations without relying on
explicit annotations. Ciga et al. [10] highlighted the benefit of amalgamating
diverse multi-organ datasets, including variations in staining and resolution,
along with an increased number of images during the SSL process for enhanced
downstream task performance. They achieved impressive results in
histopathology tasks like classification, regression, and segmentation.
Srinidhi et al. [23] introduced a domain-specific contrastive learning model
tailored for histopathology, which outperforms general-purpose contrastive
learning methods on tumor metastasis detection, tissue type classification,
and tumor cellularity quantification. They focused on enhancing
representations through a pretext task that involves predicting the order of
all feasible sequences of resolution generated from the input multi-resolution
patches. Wang et al. [25] developed a SSL approach integrating self-attention
to learn patch-level embeddings, called semantically-relevant contrastive
learning, using convolutional neural network and a multi-scale Swin
Transformer architecture as the backbone, which compares relevance between
instances to mine more positive pair. The SSL method proposed by Yang et al.
[29] comprises two self-supervised stages: cross-stain prediction and
contrastive learning, both grounded in domain-specific information. It can
merge advantages from generative and discriminative models. In addition,
SimTriplet [21] proposed by Liu et al. uses the spatial neighborhood on WSI to
provide rich positive pairs (patches with the same tissue types) for triplet
representation learning. It maximizes both intra-sample and inter-sample
similarities via triplets from positive pairs, without using negative samples.
Also, the benchmarking created by Kang et al. [15], which includes MOCO [13],
BarlowTwin [31], SwAV [6], and DINO [7], concluded that large-scale domain-
aligned pre-training is helpful for pathology, showing its value in scenarios
with limited labeled data, longer fine-tuning schedules, and when using larger
and more diverse datasets for pre-training.

Our proposed approach is related to SimTripet [21]. By leveraging a certain
number of nearby images in the self-supervised training process, we aim for a
robust patch-wise representation for achieving better results on patch-level
multi-class datasets.

##  3 Methodology

###  3.1 Overview

Inheriting the similar structure from Supervised Contrastive Learning (SupCon)
[16], our method aims to learn visual representations by multi-positive
samples. Differently from previous contrastive learning methods [8, 13]
treating one augmented view as a positive sample for the other, we harness
label information to add more positives for one sample. However, labels are
not provided in the self-supervised setting like SupCon [16]; so we adopt
nearby patches as alternative positives as illustrated in Fig. 1.

To this end, we propose a novel Nearby Patch Contrastive Learning
(NearbyPatchCL), where we maximize the similarity between not only different
views of the same patch image but also adjacent patches (See Fig. 2). Inspired
from SimTriplet [21], in our method, neighbor patches share the same tissue
(or class) because they are cropped at a small scale of WSI.

Figure 2: Overview of the proposed NearbyPatchCL. Normalized embeddings of
nearby images are pulled closer to their center while pushing away other
images.

(a) Extract center-nearby patches.

(b) Encode features and minimize loss.

Figure 3: The architecture of NearbyPatchCL includes two main parts: (a)
Extract center patches with 8 corresponding neighbors from WSIs which is a
preprocessing step discussed in section 4 and (b) A contrastive loss is used
to minimize the distance between learned features for positive samples
(center-nearby) while simultaneously maximizing the feature distance from
negative samples.

###  3.2 Nearby Patch Contrastive Learning (NearbyPatchCL)

Given a randomly sampled batch of $B$ samples (denoted as ”batch”), each image
is transformed by two random augmentations to get the training batch
containing $2B$ samples (denoted by ”multiviewed batch”). Besides, we denote
$I\equiv\\{1,\ldots,2B\\}$ as the set of indices of all samples in multiviewed
batch. For a sample with index $i$, $P(i)$ and $A(i)\equiv I\backslash P(i)$
are the collection of indices of all positives and negatives related to the
sample respectively, and $j(i)$ is the index of the other augmented sample
originating from the same source sample in the batch.

Algorithm 1 Pseudocode for NearbyPatchCL algorithm.

1:

2:Input

3: XSampled minibatch

4: CNumber of center samples in $X$

5: IIndex set of samples in the multiviewed batch

6: NNumber of nearby samples per center

7: $\mathcal{N}_{n}(x_{i})$Function return n-th nearby of sample $x_{i}$

8: $\mathcal{T}$Distribution of image transformation

9: $f$, $g$Encoder network and Projection head

10: $P(i)$Set of indices of positive samples for sample with index $i$.

11: $A(i)$Set of indices of negative samples for sample with index $i$.

12:for sampled minibatch $X=\\{x_{c}\\}_{c=1}^{C}$ do

13: $Y=\\{1,\ldots,C\\}$ # labels

14: for $n\in\\{1,\ldots,N\\}$ do# Retrieve N nearby samples

15: for $c\in\\{1\ldots C\\}$ do

16: $X\leftarrow X+\mathcal{N}_{n}(x_{c})$

17: $Y\leftarrow Y+c$

18: end for

19: end for

20: $t\sim\mathcal{T}$, $t^{\prime}\sim\mathcal{T}$ # two augmentations

21: for $k\in\\{1,\ldots,C(N+1)\\}$ do

22: $\tilde{x}_{k}=t(x_{k})$ # first augmentation

23: $h_{k}$ = $f(\tilde{x}_{k})$

24: $z_{k}=g(h_{k})$

25: $\tilde{x}_{k+C(N+1)}=t^{\prime}(x_{k})$ # second augmentation

26: $h_{k+C(N+1)}$ = $f(\tilde{x}_{k+C(N+1)})$

27: $z_{k+C(N+1)}=g(h_{k+C(N+1)})$

28: end for

29: for $i\in\\{1,\ldots,2C(N+1)\\}$ and $j\in\\{1,\ldots,2C(N+1)\\}$ do

30: $s_{i,j}=z_{i}^{T}z_{j}/\|z_{i}\|\|z_{j}\|$ # pairwise similarity

31: end for

32: for $i\in I$ do

33: $l(i)=\frac{-1}{|P(i)|}\sum\limits_{p\in
P(i)}\log\frac{\exp\left(z_{i}\cdot
z_{p}\backslash\tau\right)}{\sum\limits_{a\in A(i)}\exp\left(z_{i}\cdot
z_{a}\backslash\tau\right)}$

34: end for

35: $\mathcal{L}=\frac{1}{2C(N+1)}\sum_{i=1}^{2C(N+1)}l(i)$

36: update networks f and g to minimize $\mathcal{L}$

37:end for

38:Output Encoder $f$

Each batch contains $C$ center patches with $N$ corresponding nearby patches
($0\leq N\leq 8$), so there are $B=C\times(N+1)$ samples,
$\\{x_{k}\\}_{k=1\ldots B}$, in total. A group of patches including a center
and corresponding neighbors belongs to one class; therefore, there are $C$
classes in the batch. By applying two transformations, the multiviewed batch
contains $2C(N+1)$ samples, $\\{\widetilde{x}_{l}\\}_{l=1\ldots 2B}$. For an
arbitrary sample with index $i$, $P(i)\equiv\\{j(i)\\}\cup\mathcal{N}(i)$
where $\mathcal{N}(i)$ is the set that consists of $2N$ indices of positives
(nearby) for the sample and $2(N+1)(C-1)$ remaining instances as negatives.
Our naive loss function is as follows:

$\mathcal{L}=\sum_{i\in I}\frac{-1}{|{P}{(i)}|}\sum_{p\in
P(i)}\log\frac{\exp\left(z_{i}\cdot z_{p}/\tau\right)}{\exp\left(z_{i}\cdot
z_{p}/\tau\right)+\sum\limits_{a\in A(i)}\exp\left(z_{i}\cdot
z_{a}/\tau\right)},$ (1)

where ${z}_{i}\equiv g\left(f\left(\tilde{x}_{i}\right)\right)$ is the output
feature of sample $\tilde{x}_{i}$, the $\cdot$ symbol denotes the inner
product, and $\tau$ is a scalar temperature parameter.

Inspired by Yeh et al. [30], to address the imbalance due to cropping WSIs, we
hypothesize that robustness representations can be obtained by removing
negative samples in a batch to eliminate the negative pair contrast effect. It
means that we do not need a large batch size during the training process to
enhance learning efficiency, resulting in stable performance. This hypothesis
is explored further in Section 5.3.2. Our naive loss function (Eq. 1) becomes:

$\mathcal{L}=\sum_{i\in I}\frac{-1}{|{P}{(i)}|}\sum_{p\in
P(i)}\log\frac{\exp\left(z_{i}\cdot z_{p}/\tau\right)}{\sum\limits_{a\in
A(i)}\exp\left(z_{i}\cdot z_{a}/\tau\right)}.$ (2)

Algorithm 1 summarizes our proposed method.

###  3.3 Implementation Details

####  3.3.1 Architecture.

Figure 3 shows the overall architecture of our NearbyPatchCL, adopted from
SupCon’s architecture [16]. ResNet-50 backbone is used. The representation
from ResNet-50’s last fully connected layer remains 2048-dimensional, and it’s
then reduced to 128 dimensions using a multi-layer perceptron (MLP).

####  3.3.2 Image augmentations.

We follow the augmentation described in SupCon [16]. Additionally, horizontal
flips are randomly applied for rotation invariance in skin tissue images.
Images are resized to $128\times 128$ for computational efficiency.

####  3.3.3 Optimization.

We use the SGD optimizer, with a base learning rate of $0.2$, momentum of
$0.9$, and weight decay of $0.0001$. To optimize resources, we use Mixed
Precision Training (MPT) [19] with a learning rate scheduler. The new learning
rate is calculated as $lr\times BatchSize\times(N+1)/256$, where
$BatchSize=512/(N+1)$ is the number of center images and $N$ is the number of
nearby patches per center. The scheduler involves 10 warm-up epochs followed
by cosine decay. This SSL setup takes about 73 hours to train the encoder for
400 epochs using unlabeled data from our newly constructed P-CATCH dataset.

##  4 Proposed P-CATCH Dataset

###  4.1 Original CATCH Images

The CATCH dataset [27] consists of 350 WSIs, stored in the pyramidal Aperio
file format, having direct access to three resolution levels (0.25$\frac{\mu
m}{px}$, 1$\frac{\mu m}{px}$, 4$\frac{\mu m}{px}$). Pathologists used Slide
Runner software to create a database with 12,424 area annotations in total,
covering six non-neoplastic tissues (epidermis, dermis, subcutis, bone,
cartilage, inflammation/necrosis) and seven tumor classes. Notably, there is a
significant imbalance in the distribution of the six non-neoplastic tissue
classes within the database.

Table 1: The number of patch images in unlabeled sets, each has 247 WSIs. The
number after the dataset name denotes the number of nearby images.

Unlabeled Set | Per WSI | Total
---|---|---
Center Images | Nearby Images | Center Images | Nearby Images
P-CATCH-0 | 1,000 | 0 | 247,000 | 0
P-CATCH-1 | 500 | 500 | 123,500 | 123,500
P-CATCH-2 | 333 | 666 | 82,251 | 164,502
P-CATCH-4 | 200 | 800 | 49,400 | 197,600
P-CATCH-8 | 111 | 889 | 27,417 | 246,753

Table 2: The number of patch images per category in our P-CATCH dataset. Category | Training Set | Test Set | Unlabeled Set
---|---|---|---
Dermis | 22,020 | 81,143 | -
Epidermis | 9,471 | 16,001 | -
Inflamm/Necrosis | 19,488 | 24,612 | -
Subcutis | 16,566 | 53,426 | -
Tumor | 22,341 | 96,188 | -
Background | 1,917 | 4,533 | -
Total | 91,803 | 175,903 | 2,223,000

###  4.2 Proposed P-CATCH Dataset

####  4.2.1 Unlabeled sets.

A total of 70% of WSIs contained within the CATCH database, specifically
corresponding to 247 WSIs, is used as unlabeled data for training SSL methods.
To ensure a balanced representation of the diverse tumor subtypes present, we
adopt an equitable distribution strategy, resulting in approximately 35 WSIs
per subtype. From each WSI at the 0.25$\frac{\mu m}{px}$ resolution level, we
randomly extract image patches with size of $512\times 512$ pixels,
encompassing center patches along with nearby patches. We create different
subsets, denoted by P-CATCH-$N$, where $N$ indicates the number of nearby
samples of a center image. Table 1 shows the statistics of created five
subsets, corresponding to 0, 1, 2, 4, and 8 nearby samples. We remark that we
try to ensure an equal number of patch images for P-CATCH-$N$ subsets.

####  4.2.2 Annotated sets.

We partitioned the dataset, with the remaining 97 WSIs reserved for the test
set, while 37 WSIs from the unlabeled set were allocated for training
purposes. Specifically, we conduct random sampling to extract non-overlapping
image patches, each size at $512\times 512$ pixels, from the WSIs at the
0.25$\frac{\mu m}{px}$ resolution level. This process yields approximately
92,000 images distributed across six categories for the training set, and
around 176,000 images for the test set, spanning the same set of categories.
The distribution of images across these categories is detailed in Table 2.
This meticulous dataset partitioning ensures a rigorous and comprehensive
evaluation of our methodology.

##  5 Experiments

###  5.1 Linear Evaluation Protocol

We adhere to the well-established linear evaluation protocol [8, 31] which
entails training a linear classifier atop a frozen base network (i.e.,
ResNet-50), and test accuracy is used as a proxy for representation quality. A
single transformation is applied, involving resizing to $256\times 256$,
subsequent center-cropping to revert to the original resolution of $224\times
224$, followed by normalization. Methods are trained for 15 epochs, utilizing
SGD optimizer with a learning rate of $0.2$, momentum of $0.9$, and weight
decay of $0$. The batch size is configured at 32 for training on 1% labeled
data, while a batch size of 512 is employed for cases involving 10%, 20%, and
100% labeled data. To ensure a robust and comprehensive assessment, we employ
5-fold cross-validation on the training set, resulting in 5 trained models.
Notably, MPT is excluded during the linear evaluation. During the evaluation
of the test set, the results obtained from these 5 models are averaged to
provide a more reliable estimate of the model’s performance, thereby
effectively mitigating issues such as overfitting and yielding a more robust
evaluation outcome.

###  5.2 Comparison with State-of-the-arts

We compared our NearbyPatchCL, using $N=4$ nearby samples, with state-of-the-
art SSL methods, including SimCLR [8], SimSiam [9], BarlowTwins [31],
SimTriplet [21], and BYOL [12]. For fair evaluation, all methods employed the
ResNet-50 backbone pre-trained on the ImageNet dataset. SimTriplet [21] used
P-CATCH-1 for self-supervised training phase. Meanwhile, P-CATCH-0 subset was
used for training other SSL methods. Training parameters follow their original
work.

Table 3: Classification performance, using different amounts of training data. With different numbers of annotated data, we ensure different classes contribute similar numbers of images to address the issue that the annotation is highly imbalanced. Method | F1 Score | Balanced Accuracy
---|---|---
| 1% | 10% | 20% | 100% | 1% | 10% | 20% | 100%
Supervised baseline | 66.95 | 75.92 | 78.23 | 81.04 | 74.79 | 81.94 | 83.58 | 84.24
SimCLR [8] | 70.09 | 72.42 | 72.43 | 72.32 | 77.55 | 79.45 | 79.40 | 79.26
SimSiam [9] | 54.24 | 57.52 | 57.60 | 58.43 | 58.59 | 60.14 | 60.69 | 61.31
SimTriplet [21] | 55.18 | 57.11 | 56.20 | 58.67 | 60.75 | 62.65 | 62.78 | 63.57
BYOL [12] | 65.62 | 77.35 | 79.24 | 83.01 | 74.63 | 83.41 | 84.51 | 85.48
BarlowTwins [31] | 76.43 | 80.88 | 81.08 | 82.55 | 82.73 | 85.98 | 86.19 | 86.36
NearbyPatchCL (N=4) | 81.85 | 83.73 | 84.41 | 85.72 | 85.63 | 87.19 | 87.56 | 87.14

Comparison results in Table 3 demonstrate that our NearbyPatchCL(N=4)
significantly outperforms the supervised baseline (i.e., frozen ImageNet pre-
trained ResNet-50 with a classifier trained on the training set) and other SSL
methods across all data proportions. Furthermore, even with 1% labeled data
used for training, NearbyPatchCL still has competitive results compared to
BYOL and Barlowtwins using 100% of training data. Leveraging nearby batches as
positive samples, our approach achieves better results than state-of-the-art.
Additionally, we visualize results on a whole WSI in Fig. 4. The result also
indicates that our approach can effectively leverage unlabeled data for
improving classification tasks, making them a valuable tool in practical
scenarios where annotating WSIs is costly.

|  |
---|---|---
(a) Origin image | (b) NearPatchCL(100%) | (c) NearPatchCL(1%)

|  |
---|---|---
(c) Manual annotation | (d) Supervised(100%) | (e) Supervised(1%)

Label:TumorDermisSubcutisEpidermisInflamm-Necrosis

Figure 4: Visualization of classification results on a sample WSI from the
test set. Our proposed method shows significantly better results than that of
the supervised baseline, even with only $1\%$ of the training data.

###  5.3 Ablation Study

####  5.3.1 Number of nearby samples.

We investigate the influence of the number of nearby samples ($N$) for each
center image on the performance of the proposed method. It is notable that
with N=0, NearbyPatchCL becomes SimCLR. The result in Table 4, with DCL loss,
shows that the performance of NearbyPatchCL increases gradually from $N=0$ to
$N=4$ and achieves the best overall performance at $N=4$ in using 10%, 20%,
and 100% labeled images in the training set. The result in NearbyPatchCL(N=8)
still has a competitive result with NearbyPatchCL(N=4) and has the second-
highest overall performance. It is notable that annotations made by
pathologists may not have an accuracy of 100%. Hence, training on a large
amount of data can make the model learn some false cases, as we can see the
balanced accuracy drop when moving from $20\%$ to $100\%$ in $N=1,4,8$. On the
other hand, the performance without DCL loss increasing gradually from $N=0$
to $N=8$ also shows the superiority of utilizing nearby patches in the SSL
process.

####  5.3.2 Effect of DCL Loss.

Table 4: Ablation study with different numbers of nearby samples (N) for each image center and using DCL loss. Method | F1 Score | Balanced Accuracy
---|---|---
1% | 10% | 20% | 100% | 1% | 10% | 20% | 100%
NearbyPatchCL (N=0) | DCL | 73.83 | 78.36 | 79.14 | 80.69 | 80.01 | 83.87 | 84.35 | 84.44
NearbyPatchCL (N=1) | 77.97 | 81.84 | 82.77 | 84.40 | 83.42 | 86.54 | 87.00 | 86.94
NearbyPatchCL (N=2) | 80.27 | 82.87 | 83.43 | 85.07 | 85.63 | 87.16 | 87.38 | 87.47
NearbyPatchCL (N=4) | 81.85 | 83.73 | 84.41 | 85.72 | 85.63 | 87.19 | 87.56 | 87.14
NearbyPatchCL (N=8) | 79.66 | 83.14 | 83.97 | 85.45 | 84.60 | 87.29 | 87.53 | 87.09
NearbyPatchCL (N=0) | w/o DCL | 70.09 | 72.42 | 72.43 | 72.32 | 77.55 | 79.45 | 79.40 | 79.26
NearbyPatchCL (N=1) | 77.52 | 82.58 | 83.31 | 84.68 | 83.93 | 86.92 | 87.12 | 86.57
NearbyPatchCL (N=2) | 77.69 | 82.41 | 83.25 | 84.82 | 82.85 | 86.42 | 86.86 | 86.86
NearbyPatchCL (N=4) | 78.62 | 82.68 | 83.62 | 84.94 | 83.77 | 86.90 | 87.40 | 87.08
NearbyPatchCL (N=8) | 82.17 | 83.58 | 84.00 | 84.41 | 85.86 | 87.22 | 87.47 | 86.67

The result shown in Table 4 indicates that with more nearby samples for each
center image, as we can see in $N=\\{2,4\\}$, employing the DCL loss can
improve the performance with all percentages of labeled data.

Although our method, with DCL loss, achieves the best overall performance at
$N=4$, NearbyPatchCL (w/o DCL) still achieves the best results in 1% and 10%
of training data. With $N=1$ or $8$, NearbyPatchCL with DCL loss does not have
any significant improvement in performance, even worse at some percentages of
labeled data as in $N=1$. Our hypothesis is that NearbyPatchCL gains more
benefit from leveraging more nearby patches with DCL loss. However, with
$N=8$, there are only 111 center images that are randomly cropped from each
WSI, which are nearly half of that of $N=4$, this can make the training data
more unbalanced due to the differences in area and quantity of each category
annotation, leading to a decrease in overall performance from $N=4$ to $N=8$.

##  6 Conclusion and Future works

In this paper, we have conducted a new benchmark of patch-level multi-class
WSI classification using SSL methods. To tackle the scarcity of labeled data
and imbalanced datasets issue in digital pathology images, our work has shown
that by leveraging nearby patches as positive samples in the SSL phase, the
proposed method can have a more robust representation and perform better on
downstream tasks. Furthermore, we have shown that using DCL loss can benefit
contrastive methods while training on an imbalanced dataset.

In future work, we aim to extend our approach to other medical imaging domains
and explore its application in other downstream tasks. We also plan to further
investigate methods to enhance the interpretability of learned representations
and incorporate domain-specific knowledge to improve the performance of the
model in real-world clinical settings.

## Acknowledgement

This research was funded by Vingroup and supported by Vingroup Innovation
Foundation (VINIF) under project code VINIF.2019.DA19. This project was also
supported by the Faculty of Information Technology, University of Science,
Vietnam National University - Ho Chi Minh City.

## References

  * [1] Azizi, S., Mustafa, B., Ryan, F., Beaver, Z., Freyberg, J., Deaton, J., Loh, A., Karthikesalingam, A., Kornblith, S., Chen, T., et al.: Big self-supervised models advance medical image classification. In: ICCV. pp. 3478–3488 (2021)
  * [2] Bardes, A., Ponce, J., LeCun, Y.: Vicreg: Variance-invariance-covariance regularization for self-supervised learning. arXiv preprint arXiv:2105.04906 (2021)
  * [3] Bromley, J., Guyon, I., LeCun, Y., Säckinger, E., Shah, R.: Signature verification using a” siamese” time delay neural network. NeurIPS 6 (1993)
  * [4] Campanella, G., Hanna, M., Geneslaw, L., Miraflor, A., Silva, V., Busam, K., Brogi, E., Reuter, V., Klimstra, D., Fuchs, T.: Clinical-grade computational pathology using weakly supervised deep learning on whole slide images. Nature Medicine 25, 1 (08 2019)
  * [5] Caron, M., Bojanowski, P., Joulin, A., Douze, M.: Deep clustering for unsupervised learning of visual features. In: ECCV. pp. 132–149 (2018)
  * [6] Caron, M., Misra, I., Mairal, J., Goyal, P., Bojanowski, P., Joulin, A.: Unsupervised learning of visual features by contrasting cluster assignments. NeurIPS 33, 9912–9924 (2020)
  * [7] Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., Joulin, A.: Emerging properties in self-supervised vision transformers. In: ICCV. pp. 9650–9660 (2021)
  * [8] Chen, T., Kornblith, S., Norouzi, M., Hinton, G.: A simple framework for contrastive learning of visual representations. In: ICML. pp. 1597–1607 (2020)
  * [9] Chen, X., He, K.: Exploring simple siamese representation learning. In: CVPR. pp. 15750–15758 (2021)
  * [10] Ciga, O., Xu, T., Martel, A.L.: Self supervised contrastive learning for digital histopathology. Machine Learning with Applications 7, 100198 (2022)
  * [11] Gamper, J., Rajpoot, N.: Multiple instance captioning: Learning representations from histopathology textbooks and articles. In: CVPR. pp. 16549–16559 (2021)
  * [12] Grill, J.B., Strub, F., Altché, F., Tallec, C., Richemond, P., Buchatskaya, E., Doersch, C., Avila Pires, B., Guo, Z., Gheshlaghi Azar, M., et al.: Bootstrap your own latent-a new approach to self-supervised learning. NeurIPS 33, 21271–21284 (2020)
  * [13] He, K., Fan, H., Wu, Y., Xie, S., Girshick, R.: Momentum contrast for unsupervised visual representation learning. CVPR pp. 9729–9738 (2020)
  * [14] Jiang, Z., Chen, T., Mortazavi, B.J., Wang, Z.: Self-damaging contrastive learning. In: ICML. pp. 4927–4939 (2021)
  * [15] Kang, M., Song, H., Park, S., Yoo, D., Pereira, S.: Benchmarking self-supervised learning on diverse pathology datasets. In: CVPR. pp. 3344–3354 (2023)
  * [16] Khosla, P., Teterwak, P., Wang, C., Sarna, A., Tian, Y., Isola, P., Maschinot, A., Liu, C., Krishnan, D.: Supervised contrastive learning. NeurIPS 33, 18661–18673 (2020)
  * [17] van der Laak, J.A., Litjens, G.J.S., Ciompi, F.: Deep learning in histopathology: the path to the clinic. Nature Medicine 27, 775 – 784 (2021)
  * [18] Li, B., Li, Y., Eliceiri, K.W.: Dual-stream multiple instance learning network for whole slide image classification with self-supervised contrastive learning. In: CVPR. pp. 14318–14328 (2021)
  * [19] Li, H., Wang, Y., Hong, Y., Li, F., Ji, X.: Layered mixed-precision training: A new training method for large-scale ai models. Journal of King Saud University - Computer and Information Sciences 35(8), 101656 (2023)
  * [20] Liu, K., Zhu, W., Shen, Y., Liu, S., Razavian, N., Geras, K.J., Fernandez-Granda, C.: Multiple instance learning via iterative self-paced supervised contrastive learning. In: CVPR. pp. 3355–3365 (2023)
  * [21] Liu, Q., Louis, P.C., Lu, Y., Jha, A., Zhao, M., Deng, R., Yao, T., Roland, J.T., Yang, H., Zhao, S., et al.: Simtriplet: Simple triplet representation learning with a single gpu. In: MICCAI. pp. 102–112 (2021)
  * [22] Lu, M.Y., Williamson, D.F., Chen, T.Y., Chen, R.J., Barbieri, M., Mahmood, F.: Data-efficient and weakly supervised computational pathology on whole-slide images. Nature biomedical engineering 5(6), 555–570 (2021)
  * [23] Srinidhi, C.L., Kim, S.W., Chen, F.D., Martel, A.L.: Self-supervised driven consistency training for annotation efficient histopathology image analysis. Medical Image Analysis 75, 102256 (2022)
  * [24] Wang, D., Khosla, A., Gargeya, R., Irshad, H., Beck, A.H.: Deep learning for identifying metastatic breast cancer. arXiv preprint arXiv:1606.05718 (2016)
  * [25] Wang, X., Yang, S., Zhang, J., Wang, M., Zhang, J., Huang, J., Yang, W., Han, X.: Transpath: Transformer-based self-supervised learning for histopathological image classification. In: MICCAI. pp. 186–195 (2021)
  * [26] Wang, X., Yang, S., Zhang, J., Wang, M., Zhang, J., Yang, W., Huang, J., Han, X.: Transformer-based unsupervised contrastive learning for histopathological image classification. Medical image analysis 81, 102559 (2022)
  * [27] Wilm, F., Fragoso, M., Marzahl, C., Bertram, C., Klopfleisch, R., Maier, A., Aubreville, M., Breininger, K.: Canine cutaneous cancer histology dataset (2022), https://wiki.cancerimagingarchive.net/x/DYITBg
  * [28] Wulczyn, E., Steiner, D., Moran, M., Plass, M., Reihs, R., Tan, F., Flament-Auvigne, I., Brown, T., Regitnig, P., Chen, P.H., Hegde, N., Sadhwani, A., MacDonald, R., Ayalew, B., Corrado, G., Peng, L., Tse, D., Müller, H., Xu, Z., Mermel, C.: Interpretable survival prediction for colorectal cancer using deep learning. npj Digital Medicine 4, 71 (04 2021)
  * [29] Yang, P., Yin, X., Lu, H., Hu, Z., Zhang, X., Jiang, R., Lv, H.: Cs-co: A hybrid self-supervised visual representation learning method for h&e-stained histopathological images. Medical Image Analysis 81, 102539 (2022)
  * [30] Yeh, C.H., Hong, C.Y., Hsu, Y.C., Liu, T.L., Chen, Y., LeCun, Y.: Decoupled contrastive learning. In: ECCV. pp. 668–684 (2022)
  * [31] Zbontar, J., Jing, L., Misra, I., LeCun, Y., Deny, S.: Barlow twins: Self-supervised learning via redundancy reduction. In: ICML. pp. 12310–12320 (2021)
  * [32] Zheng, Y., Jiang, B., Shi, J., Zhang, H., ying Xie, F.: Encoding histopathological wsis using gnn for scalable diagnostically relevant regions retrieval. In: MICCAI (2019)