0  marginparsep has been altered.
topmargin has been altered.
marginparwidth has been altered.
marginparpush has been altered.
The page layout violates the ICML style. Please do not change the page layout,
or include packages like geometry, savetrees, or fullpage, which change it for
you. We’re not able to reliably undo arbitrary changes to the style. Please
remove the offending package(s), or layout-changing commands and try again.

Checkmate: Breaking the Memory Wall
with Optimal Tensor Rematerialization

Anonymous Authors1

###### Abstract

We formalize the problem of trading-off DNN training time and memory
requirements as the tensor rematerialization optimization problem, a
generalization of prior checkpointing strategies. We introduce Checkmate, a
system that solves for optimal rematerialization schedules in reasonable times
(under an hour) using off-the-shelf MILP solvers or near-optimal schedules
with an approximation algorithm, then uses these schedules to accelerate
millions of training iterations. Our method scales to complex, realistic
architectures and is hardware-aware through the use of accelerator-specific,
profile-based cost models. In addition to reducing training cost, Checkmate
enables real-world networks to be trained with up to 5.1$\times$ larger input
sizes. Checkmate is an open-source project, available at
https://github.com/parasj/checkmate.

††footnotetext:  1Anonymous Institution, Anonymous City, Anonymous Region,
Anonymous Country.
Correspondence to: Anonymous Author<EMAIL_ADDRESS>
Preliminary work. Under review by the Machine Learning and Systems (MLSys)
Conference. Do not distribute.

##  1 Introduction

Deep learning training workloads demand large amounts of high bandwidth
memory. Researchers are pushing the memory capacity limits of hardware
accelerators such as GPUs by training neural networks on high-resolution
images Dong et al. (2016); Kim et al. (2016); Tai et al. (2017), 3D point-
clouds Chen et al. (2017); Yang et al. (2018), and long natural language
sequences Vaswani et al. (2017); Devlin et al. (2018); Child et al. (2019). In
these applications, training memory usage is dominated by the intermediate
activation tensors needed for backpropagation (Figure 3).

The limited availability of high bandwidth on-device memory creates a memory
wall that stifles exploration of novel architectures. Across applications,
authors of state-of-the-art models cite memory as a limiting factor in deep
neural network (DNN) design Krizhevsky et al. (2012); He et al. (2016); Chen
et al. (2016a); Gomez et al. (2017); Pohlen et al. (2017); Child et al.
(2019); Liu et al. (2019); Dai et al. (2019).

As there is insufficient RAM to cache all activation tensors for
backpropagation, some select tensors can be discarded during forward
evaluation. When a discarded tensor is necessary as a dependency for gradient
calculation, the tensor can be rematerialized. As illustrated in Figure 1,
rematerializing values allows a large DNN to fit within memory at the expense
of additional computation.

Figure 1: This 32-layer deep neural network requires 30GB of memory during
training in order to cache forward pass activations for the backward pass.
Freeing certain activations early and rematerializing them later reduces
memory requirements by 21GB at the cost of a modest runtime increase.
Rematerialized layers are denoted as shaded blue regions. We present
Checkmate, a system to rematerialize large neural networks _optimally_.
Checkmate is hardware-aware, memory-aware and supports arbitrary DAGs.

Griewank & Walther (2000) and Chen et al. (2016b) present heuristics for
rematerialization when the forward pass forms a linear graph, or path graph.
They refer to the problem as checkpointing. However, their approaches cannot
be applied generally to nonlinear DNN structures such as residual connections,
and rely on the strong assumption that all nodes in the graph have the same
cost. Prior work also assumes that gradients may never be rematerialized.
These assumptions limit the efficiency and generality of prior approaches.

Our work formalizes tensor rematerialization as a constrained optimization
problem. Using off-the-shelf numerical solvers, we are able to discover
optimal rematerialization strategies for arbitrary deep neural networks in
TensorFlow with non-uniform computation and memory costs. We demonstrate that
optimal rematerialization allows larger batch sizes and substantially reduced
memory usage with minimal computational overhead across a range of image
classification and semantic segmentation architectures. As a consequence, our
approach allows researchers to easily explore larger models, at larger batch
sizes, on more complex signals with minimal computation overhead.

Figure 2: Overview of the Checkmate system.

In particular, the contributions of this work include:

  * •

a formalization of the rematerialization problem as a mixed integer linear
program with a substantially more flexible search space than prior work, in
Section 4.7.

  * •

a fast approximation algorithm based on two-phase deterministic LP rounding,
in Section 5.

  * •

Checkmate, a system implemented in TensorFlow that enables training models
with up to 5.1$\times$ larger input sizes than prior art at minimal overhead.

##  2 Motivation

While inference optimizations are well studied Jain et al. (2019), training
workloads have received less attention. Memory consumption during training
consists of (a) intermediate features, or activations, whose size depends on
input dimensions and (b) parameters and their gradients whose size depends on
weight dimensions. Given that inputs are often several order of magnitude
larger than kernels, most memory is used by features, demonstrated in Figure
3.

Frameworks such as TensorFlow Abadi et al. (2016) and PyTorch Paszke et al.
(2017; 2019) store all activations during the forward pass. Gradients are
backpropagated from the loss node, and each activation is freed after its
gradient has been calculated. In Figure 1, we compare this memory intensive
policy and a rematerialization strategy for a real neural network. Memory
usage is significantly reduced by deallocating some activations in the forward
pass and recomputing them in the backward pass. Our goal is fit an arbitrary
network within our memory budget while incurring the minimal additional
runtime penalty from recomputation.

Most prior work assumes networks have linear graphs. For example, Chen et al.
(2016b) divides the computation into $\sqrt{n}$ segments, each with $\sqrt{n}$
nodes. Each segment endpoint is stored during the forward pass. During the
backward pass, segments are recomputed in reverse order at $O(n)$ cost.

Linear graph assumptions limit applicability of prior work. For example, while
the popular ResNet50 He et al. (2016) can be linearized by treating each
residual block as a single node, this leads to inefficient solutions. For
networks with longer skip connections, e.g., U-Net Ronneberger et al. (2015),
grouping nodes oversimplifies the graph.

Prior work also assumes all layers are equally expensive to recompute. In the
VGG19 architecture Simonyan & Zisserman (2014), the largest layer is _six
orders of magnitude_ more expensive than the smallest layer.

Our work makes few assumptions on neural network graphs. We explore a solution
space that allows for (a) arbitrary graphs with several inputs and outputs for
each node, (b) variable memory costs across layers and (c) variable
computation costs for each layer (such as FLOPs or profiled runtimes). We
constrain solutions to simply be _correct_ (a node’s dependencies must be
materialized before it can be evaluated) and _within the RAM budget_ (at any
point during execution, resident tensors must fit into RAM).

Figure 3: Memory consumed by activations far outweigh parameters for popular
model architectures. Moreover, advances in GPU DRAM capacity are quickly
utilized by researchers; the dashed line notes the memory limit of the GPU
used to train each model.

Subject to these constraints, we find solutions that minimize the amount of
time it takes to perform a single training iteration. We project schedules
into space and time, allowing us to cast the objective as a linear expression.
This problem can then be solved using off-the-shelf mixed integer linear
program solvers such as GLPK  or COIN-OR Branch-and-Cut Forrest et al. (2019).
An optimal solution to the MILP will minimize the amount of additional compute
cost within the memory budget.

##  3 Related Work

Method |  Description |  General graphs |  Cost aware |  Memory aware
---|---|---|---|---
Checkpoint all (Ideal) |  No rematerialization. Default in deep learning frameworks. |  $\surd$ |  $\times$ |  $\times$
Griewank et al. $\log{n}$ |  Griewank & Walther (2000) revolve procedure |  $\times$ |  $\times$ |  $\times$
Chen et al. $\sqrt{n}$ |  Chen et al. (2016b) checkpointing heuristic |  $\times$ |  $\times$ |  $\times$
Chen et al. greedy |  Chen et al. (2016b), with search over parameter $b$ |  $\times$ |  $\times$ |  $\sim$
AP $\sqrt{n}$ |  Chen et al. $\sqrt{n}$ on articulation points + optimal R solve |  $\sim$ |  $\times$ |  $\times$
AP greedy |  Chen et al. greedy on articulation points + optimal R solve |  $\sim$ |  $\times$ |  $\sim$
Linearized $\sqrt{n}$ |  Chen et al. $\sqrt{n}$ on topological sort + optimal R solve |  $\surd$ |  $\times$ |  $\times$
Linearized greedy |  Chen et al. greedy on topological sort + optimal R solve |  $\surd$ |  $\times$ |  $\sim$
Checkmate ILP |  Our ILP as formulated in Section 4 |  $\surd$ |  $\surd$ |  $\surd$
Checkmate approx. |  Our LP rounding approximation algorithm (Section 5) |  $\surd$ |  $\surd$ |  $\surd$

Table 1: Rematerialization baselines and our extensions to make them
applicable to non-linear architectures

We categorize related work as checkpointing, reversible networks, distributed
computation, and activation compression.

Checkpointing and rematerialization Chen et al. (2016b) propose a heuristic
for checkpointing idealized unit-cost linear $n$-layer graphs with
$O(\sqrt{n})$ memory usage. Griewank & Walther (2000) checkpoint similar
linear unit-cost graphs with $O(\log{n})$ memory usage and prove optimality
for linear graphs with unit per-node cost and memory. In practice, DNN layers
vary significantly in memory usage and computational cost Sze et al. (2017),
so these heuristics are not optimal in practice. Chen et al. (2016b) also
develop a greedy algorithm that checkpoints layers of a network in roughly
memory equal segments, with a hyperparameter $b$ for the size of such
segments. Still, neither procedure is cost-aware nor deallocates checkpoints
when possible. Gruslys et al. (2016) develop a dynamic programming algorithm
for checkpoint selection in unrolled recurrent neural network training,
exploiting their linear forward graphs. Feng & Huang (2018) provide a dynamic
program to select checkpoints that partition branching networks, but ignore
layer costs and memory usage. Siskind & Pearlmutter (2018a) develop a divide-
and-conquer strategy in programs. Beaumont et al. (2019) use dynamic
programming for checkpoint selection in a specific architecture with joining
sub-networks.

Intermediate value recomputation is also common in register allocation.
Compiler backends lower an intermediate representation of code to an
architecture-specific executable binary. During lowering, an abstract static
single assignment (SSA) graph of values and operations Rosen et al. (1988);
Cytron et al. (1991) is concretized by mapping values to a finite number of
registers. If insufficient registers are available for an SSA form computation
graph, values are spilled to main memory by storing and later loading the
value. Register allocation has been formulated as graph coloring problem
Chaitin et al. (1981), integer program Goodwin & Wilken (1996); Lozano et al.
(2018), and network flow Koes & Goldstein (2006).

Register allocators may recompute constants and values with register-resident
dependencies if the cost of doing so is less than the cost of a spill Chaitin
et al. (1981); Briggs et al. (1992); Punjani (2004). While similar to our
setup, register rematerialization is limited to exceptional values that can be
recomputed in a single instruction with dependencies already in registers. For
example, memory offset computations can be cheaply recomputed, and loads of
constants can be statically resolved. In contrast, Checkmate can recompute
entire subgraphs of the program’s data-flow.

During the evaluation of a single kernel, GPUs spill per-thread registers to a
thread-local region of global memory (i.e. local memory) Micikevicius (2011);
NVIDIA (2017). NN training executes DAGs of kernels and stores intermediate
values in shared global memory. This produces a high range of value sizes,
from 4 byte floats to gigabyte tensors, whereas CPU and GPU registers range
from 1 to 64 bytes. Our problem of interkernel memory scheduling thus differs
in scale from the classical problem of register allocation within a kernel or
program. Rematerialization is more appropriate than copying values out of core
as the cost of spilling values from global GPU memory to main memory (RAM) is
substantial Micikevicius (2011); Jain et al. (2018), though possible Meng et
al. (2017).

Reversible Networks Gomez et al. (2017) propose a reversible (approximately
invertible) residual DNN architecture, where intermediate temporary values can
be recomputed from values derived later in the standard forward computation.
Reversibility enables recomputation during the backward pass. Bulo et al.
(2018) replace only ReLU and batch normalization layers with invertible
variants and reduce memory usage up to 50%. We find rematerialization enables
greater savings and a wider range of budgets, but reversibility is a promising
complementary approach.

Distributed computation Orthogonal approaches to address the limited memory
problem are distributed-memory computations and gradient accumulation.
However, model parallelism requires access to additional expensive compute
accelerators, fast networks, and non-trivial partitioning of model state to
balance communication and computation Gholami et al. (2018); Jia et al.
(2018b); McCandlish et al. . Gradient accumulation enables larger batch sizes
by computing the gradients with multiple sub-batches across a mini-batch.
However, gradient accumulation can degrade performance as batch normalization
performs poorly on small batch sizes Wu & He (2018); Ioffe & Szegedy (2015).

Activation compression In some DNN applications, it is possible to process
compressed representations with minimal accuracy loss. Gueguen et al. (2018)
classify discrete cosine transforms of JPEG images rather than raw images.
Jain et al. (2018) quantize activations, cutting memory usage in half.
Compression reduces memory usage by a constant factor, but reduces accuracy.
Our approach is mathematically equivalent to rematerialization-free training
and incurs no accuracy penalty.

##  4 Optimal Rematerialization

In this section, we develop an optimal solver that schedules computation and
garbage collection during the evaluation of general data-flow graphs including
those used in neural network training. Our proposed scheduler minimizes
computation or execution time while guaranteeing that the schedule will not
exceed device memory limitations. The rematerialization problem is formulated
as a mixed integer linear program (MILP) that can be solved with standard
commercial or open-source solvers.

###  4.1 Problem definition

A computation or data-flow graph ${G=(V,E)}$ is a directed acyclic graph with
$n$ nodes ${V=\\{v_{1},\ldots,v_{n}\\}}$ that represent operations yielding
values (e.g. tensors). Edges represent dependencies between operators, such as
layer inputs in a neural network. Nodes are numbered according to a
topological order, such that operation $v_{j}$ may only depend on the results
of operations $v_{i<j}$.

Each operator’s output takes $M_{v}$ memory to store and costs $C_{v}$ to
compute from its inputs. We wish to find the terminal node $v_{n}$ with peak
memory consumption under a memory budget, $M_{\text{budget}}$, and minimum
total cost of computation.

Figure 4: Dependencies of $v_{k}$ can only be garbage collected after it is
evaluated. $U_{t,k}$ measures the memory used after evaluating $v_{k}$ and
before deallocating its dependencies. $v_{b}$ and $v_{c}$ may be deallocated
during garbage collection, but $v_{a}$ may not due to a forward edge.

###  4.2 Representing a schedule

We represent a schedule as a series of nodes being saved or (re)computed. We
unroll the execution of the network into T stages and only allow a node to be
computed once per stage. $S_{t,i}\in\\{0,1\\}$ indicates that the result of
operation $i$ should be retained in memory at stage $t-1$ until stage $t$. We
also define $R_{t,i}\in\\{0,1\\}$ be a binary variable reflecting whether
operation i is recomputed at time step t.

Our representation generalizes checkpointing Griewank & Walther (2000); Chen
et al. (2016b); Gruslys et al. (2016); Siskind & Pearlmutter (2018b); Feng &
Huang (2018), as values can be retained and deallocated many times, but comes
at the cost of $O(Tn)$ decision variables.

To trade-off the number of decision variables and schedule flexibility, we
limit $T$ to $T=n$. This allows for $O(n^{2})$ operations and constant memory
in linear graphs.

###  4.3 Scheduling with ample memory

First, consider neural network evaluation on a processor with ample memory.
Even without a memory constraint, our solver must ensure that checkpointed and
computed operations have dependencies resident in memory. Minimizing the total
cost of computation across stages with dependency constraints yields objective
(4.3): —l—[1] R,S∑_t=1^n ∑_i=1^t C_i R_t,i R_t, j≤R_t, i + S_t, i∀t ∀(v_i,
v_j) ∈E S_t, i≤R_t-1, i + S_t-1, i ∀t ≥2 ∀i ∑i S1, i= 0 ∑t Rt,n≥1 R_t, i, S_t,
i ∈{0, 1}∀t ∀i Constraints ensure feasibility and completion. Constraint (4.3)
and (4.3) ensure that an operation is computed in stage $t$ only if all
dependencies are available. To cover the edge case of the first stage,
constraint (4.3) specifies that no values are initially in memory. Finally,
covering constraint (4.3) ensures that the last node in the topological order
is computed at some point in the schedule so that training progresses.

###  4.4 Constraining memory utilization

To constrain memory usage, we introduce memory accounting variables
${U_{t,k}\in\mathbb{R}_{+}}$ into the ILP. Let $U_{t,k}$ denote the memory
used just after computing node $v_{k}$ in stage $t$. $U_{t,k}$ is defined
recursively in terms of auxiliary binary variables $\textsc{Free}_{t,i,k}$ for
$(v_{i},v_{k})\in E$, which specifies whether node $v_{i}$ may be deallocated
in stage $t$ after evaluating node $v_{k}$.

We assume that (1) network inputs and parameters are always resident in memory
and (2) enough space is allocated for gradients of the loss with respect to
parameters.111While gradients can be deleted after updating parameters, we
reserve constant space since many parameter optimizers such as SGD with
momentum maintain gradient statistics. Parameter gradients are typically
small, the same size as the parameters themselves. Additionally, at the
beginning of a stage, all checkpointed values are resident in memory. Hence,
we initialize the recurrence,

$U_{t,0}=\underbrace{M_{\text{input}}+2M_{\text{param}}}_{\text{Constant
overhead}}+\sum_{i=1}^{n}\underbrace{M_{i}S_{t,i}}_{\text{Checkpoints}}$ (1)

Suppose $U_{t,k}$ bytes of memory are in use after evaluating $v_{k}$. Before
evaluating $v_{k+1}$, $v_{k}$ and dependencies (parents) of $v_{k}$ may be
deallocated if there are no future uses. Then, an output tensor for the result
of $v_{k+1}$ is allocated, consuming memory $M_{k+1}$. The timeline is
depicted in Figure 4, yielding recurrence (2):

$U_{t,k+1}=U_{t,k}-\text{mem\\_freed}_{t}(v_{k})+R_{t,k+1}M_{k+1},$ (2)

where $\text{mem\\_freed}_{t}(v_{k})$ is the amount of memory freed by
deallocating $v_{k}$ and its parents at stage $t$. Let

$\displaystyle\textsc{Deps}[k]$ $\displaystyle=\\{i:(v_{i},v_{k})\in
E\\},\text{ and}$ $\displaystyle\textsc{Users}[i]$
$\displaystyle=\\{j:(v_{i},v_{j})\in E\\}$

denote parents and children of a node, respectively. Then, in terms of
auxiliary variable $\textsc{Free}_{t,i,k}$, for $(v_{i},v_{k})\in E$,

$\displaystyle\text{mem\\_freed}_{t}(v_{k})=\textstyle\sum_{\begin{subarray}{c}i\in\textsc{Deps}[k]\\\
\cup\\{k\\}\end{subarray}}M_{i}*\textsc{Free}_{t,i,k},\leavevmode\nobreak\
\text{and}$ (3)
$\displaystyle\textsc{Free}_{t,i,k}=R_{t,k}*\underbrace{(1-S_{t+1,i})}_{\text{Not
checkpoint}}\prod_{\begin{subarray}{c}j\in\textsc{Users}[i]\\\
j>k\end{subarray}}\underbrace{(1-R_{t,j})}_{\text{Not dep.}}$ (4)

The second factor in (4) ensures that $M_{i}$ bytes are freed only if $v_{i}$
is not checkpointed for the next stage. The final factors ensure that
$\textsc{Free}_{t,i,k}=0$ if any child of $v_{i}$ is computed in the stage,
since then $v_{i}$ needs to be retained for later use. Multiplying by
$R_{t,k}$ in (4) ensures that values are only freed at most once per stage
according to Theorem 4.1,

######  Theorem 4.1 (No double deallocation).

If (4) holds for all $(v_{i},v_{k})\in E$, then
$\sum_{k\in\textsc{Users}[i]}\textsc{Free}_{t,i,k}\leq 1\leavevmode\nobreak\
\forall t,i$.

###### Proof.

Assume for the sake of contradiction that $\exists
k_{1},k_{2}\in\textsc{Users}[i]$ such that
$\textsc{Free}_{t,i,k_{1}}=\textsc{Free}_{t,i,k_{2}}=1$. By the first factor
in (4), we must have $R_{t,k_{1}}=R_{t,k_{2}}=1$. Assume without loss of
generality that $k_{2}>k_{1}$. By the final factor in (4), we have
$\textsc{Free}_{t,i,k_{1}}\leq 1-R_{t,k_{2}}=0$, which is a contradiction. ∎

###  4.5 Linear reformulation of memory constraint

While the recurrence (1-2) defining $U$ is linear, the right hand size of (4)
is a polynomial. To express Free in our ILP, it must be defined via linear
constraints. We rely on Lemma 4.1 and 4.2 to reformulate (4) into a tractable
form.

######  Lemma 4.1 (Linear Reformulation of Binary Polynomial).

If $x_{1},\ldots,x_{n}\in\\{0,1\\}$, then

$\prod_{i=1}^{n}x_{i}=\begin{cases}1&\sum_{i=1}^{n}(1-x_{i})=0\\\
0&\text{otherwise}\end{cases}$

###### Proof.

If all $x_{1},\ldots,x_{n}=1$, then $\sum_{i=1}^{n}(1-x_{i})=0$ and we have
$\Pi_{i=1}^{n}x_{i}=1$. If otherwise any $x_{j}=0$, then we have
$\Pi_{i=1}^{n}x_{i}=0$, as desired. This can also be seen as an application of
De Morgan’s laws for boolean arithmetic. ∎

######  Lemma 4.2 (Linear Reformulation of Indicator Constraints).

Given $0\leq y\leq\kappa$ where $y$ is integral and $\kappa$ is a constant
upper bound on $y$, then

$\displaystyle x=\begin{cases}1&y=0\\\ 0&\text{otherwise}\end{cases}$

if and only if $x\in\\{0,1\\}$ and $(1-x)\leq y\leq\kappa(1-x)$.

###### Proof.

For the forward direction, first note that by construction, $x\in\\{0,1\\}$.
If $y=0$ and $x=1$, then $(1-x)=0\leq y\leq 0=\kappa(1-x)$. Similarly, if
$y\geq 1$ and $x=0$, then $1\leq y\leq\kappa$, which is true since $0\leq
y\leq\kappa$ and $y$ is integral. The converse holds similarly. ∎

To reformulate Constraint 4, let $\text{num\\_hazards}(t,i,k)$ be the number
of zero factors on the RHS of the constraint. This is a linear function of the
decision variables,

$\text{num\\_hazards}(t,i,k)=(1-R_{t,k})+S_{t+1,i}+\sum_{\begin{subarray}{c}j\in\textsc{Users}[i]\\\
j>k\end{subarray}}R_{t,j}$

Applying Lemma 4.1 to the polynomial constraint, we have,

$\displaystyle\textsc{Free}_{t,i,k}=\begin{cases}1&\text{num\\_hazards}(t,i,k)=0\\\
0&\text{otherwise}\end{cases}$ (5)

By Lemma 4.2, if $\kappa$ is the maximum value that
$\text{num\\_hazards}(t,i,k)$ can assume, the following constraints are
equivalent to (5),

$\displaystyle\textsc{Free}_{t,i,k}$ $\displaystyle\in\\{0,1\\}$ (6a)
$\displaystyle 1-\textsc{Free}_{t,i,k}$
$\displaystyle\leq\text{num\\_hazards}(t,i,k)$ (6b)
$\displaystyle\kappa(1-\textsc{Free}_{t,i,k})$
$\displaystyle\geq\text{num\\_hazards}(t,i,k)$ (6c)

###  4.6 Tractability via frontier-advancing stages

Fixing the execution order of nodes in the graph can improve the running time
of the algorithm. In eager-execution frameworks such as PyTorch, the order is
given by user code and operations are executed serially. Separating ordering
and allocation is common in compiler design, and both LLVM Lattner (2002) and
GCC Olesen (2011) have separate instruction scheduling and register allocation
passes.

Any topological order of the nodes is a possible execution order. Given a
topological order, such as the one introduced in Section 4.1, we partition the
schedule into frontier-advancing stages such that node $v_{i}$ is evaluated
for the first time in stage $i$. We replace constraints (4.3, 4.3) that ensure
the last node is computed with stricter constraints (7a-7c),

$\displaystyle R_{i,i}$ $\displaystyle=1\leavevmode\nobreak\
\leavevmode\nobreak\ \forall i\leavevmode\nobreak\ \leavevmode\nobreak\
\leavevmode\nobreak\ \text{(frontier-advancing partitions)}$ (7a)
$\displaystyle\textstyle\sum_{i\geq t}S_{t,i}$
$\displaystyle=0\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
\text{(lower tri., no initial checkpoints)}$ (7b)
$\displaystyle\textstyle\sum_{i>t}R_{t,i}$
$\displaystyle=0\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
\text{(lower triangular)}$ (7c)

This reduces the feasible set, constraining the search space and improving
running time. For an 8 layer ($n=17$) linear graph neural network with unit
$C_{i},M_{i}$ at a memory budget of 4, Gurobi optimizes the unpartitioned MILP
in 9.4 hours and the partitioned MILP in 0.23 seconds to the same objective.
In Appendix A, we analyze the integrality gap of both forms of the problem to
understand the speedup.

###  4.7 Complete Integer Linear Program formulation

The complete memory constrained MILP follows in (4.7), with $O(|V||E|)$
variables and constraints. —l—[0] R,S,U,Free_∑_t=1^n ∑_i=1^t C_i R_t,i (4.3),
(4.3), (4.3), (1), (2) (6a), (6b), (6c), (7a), (7b), (7c) U_t,k ≤M_budget

###  4.8 Constraints implied by optimality

Problem 4.7 can be simplified by removing constraints implied by optimality of
a solution. ${\textsc{Free}_{t,k,k}=1}$ only if operation $k$ is spuriously
evaluated with no uses of the result. Hence, the solver can set $R_{t,k}=0$ to
reduce cost. We eliminate $|V|^{2}$ variables $\textsc{Free}_{t,k,k}$, assumed
to be 0, by modifying (3) to only sum over $i\in\textsc{Deps}[k]$. These
variables can be computed inexpensively after solving.

###  4.9 Generating an execution plan

Given a feasible solution to (4.7), $(R,S,U,\textsc{Free})$, Algorithm 1
generates an execution plan via a row major scan of $R$ and $S$ with
deallocations determined by Free. An execution plan is a program
$P=(s_{1},\ldots,s_{k})$ with $k$ statements. When statement %r = compute v is
interpreted, operation $v$ is evaluated. The symbol %r denotes a virtual
register used to track the resulting value. Statement deallocate %r marks the
value tracked by virtual register %r for garbage collection.

The execution plan generated by Algorithm 1 is further optimized by moving
deallocations earlier in the plan when possible. Spurious checkpoints that are
unused in a stage can be deallocated at the start of the stage rather than
during the stage. Still, this code motion is unnecessary for feasibility as
the solver guarantees that the unoptimized schedule will not exceed the
desired memory budget.

The execution plan can either be interpreted during training, or encoded as a
static computation graph. In this work, we generate a static graph
$G^{\prime}=(V^{\prime},E^{\prime})$ from the plan, which is executed by a
numerical machine learning framework. See Section 6.2 for implementation
details.

Algorithm 1 Generate execution plan

Input: graph $G=(V,E)$, feasible $(R,S,\textsc{Free})$

Output: execution plan $P=(s_{1},\ldots,s_{k})$

Initialize $\textsc{Regs}[1\ldots|V|]=-1$, $r=0$, $P=()$.

for $t=1$ to $|V|$ do

for $k=1$ to $|V|$ do

if $R_{t,k}$ then

// Materialize $v_{k}$

add %$r$ = compute vk to $P$

$\textsc{Regs}[k]=r$

$r=r+1$

end if

// Free $v_{k}$ and dependencies

for $i\in\textsc{Deps}[k]\cup\\{k\\}$ do

if $\textsc{Free}_{t,i,k}$ then

add deallocate %$\textsc{Regs}[i]$ to $P$

end if

end for

end for

end for

return $P$

###  4.10 Cost model

To estimate the runtime of a training iteration under a rematerialization
plan, we apply an additive cost model (4.3), incurring cost $C_{i}$ when node
$v_{i}$ is evaluated. Costs are determined prior to MILP construction by
profiling network layers on target hardware with random inputs across a range
of batch sizes and input shapes, and exclude static graph construction and
input generation time. As neural network operations consist of dense numerical
kernels such as matrix multiplication, these runtimes are low variance and
largely independent of the specific input data Jia et al. (2018a); Sivathanu
et al. (2019). However, forward pass time per batch item decreases with
increasing batch size due to improved data parallelism Canziani et al. (2016),
so it is important to compute costs with appropriate input dimensions.

The memory consumption of each value in the data-flow graph is computed
statically as input and output sizes are known. Values are dense, multi-
dimensional tensors stored at 4 byte floating point precision. The computed
consumption $M_{i}$ is used to construct memory constraints (1-2).

##  5 Approximation

Many of our benchmark problem instances are tractable to solve using off-the-
shelf integer linear program solvers, with practical solve times ranging from
seconds to an hour. ILP results in this paper are obtained with a 1 hour time
limit on a computer with at least 24 cores. Relative to training time, e.g. 21
days for the BERT model Devlin et al. (2018), solving the ILP adds less than a
percent of runtime overhead.

While COTS solvers such as COIN-OR Forrest et al. (2019) leverage methods like
branch-and-bound to aggressively prune the decision space, they can take
superpolynomial time in the worst-case and solving ILPs is NP-hard in general.
In the worst-case, for neural network architectures with hundreds of layers,
it is not feasible to solve the rematerialization problem via our ILP. An
instance of the VGG16 architecture Simonyan & Zisserman (2014) takes seconds
to solve. For DenseNet161 Huang et al. (2017), no feasible solution was found
within one day.

For many classical NP-hard problems, approximation algorithms give solutions
close to optimal with polynomial runtime. We review a linear program that
produces fractional solutions in polynomial time in Section 5.1. Using the
fractional solutions, we present a two-phase rounding algorithm in Section 5.2
that rounds a subset of the decision variables, then finds a minimum cost,
feasible setting of the remaining variables to find near-optimal integral
solutions.

###  5.1 Relaxing integrality constraints

By relaxing integrality constraints (4.3), the problem becomes trivial to
solve as it is a linear program over continuous variables. It is well known
that an LP is solvable in polynomial time via Karmarkar’s algorithm Karmarkar
(1984) or barrier methods Nesterov & Nemirovskii (1994). With relaxation
$R,S,\textsc{Free}\in[0,1]$, the objective (4.3) defines a lower-bound for the
cost of the optimal integral solution.

Rounding is a common approach to find approximate integral solutions given the
result of an LP relaxation. For example, one can achieve a
$\frac{3}{4}$-approximation for MAX SAT Yannakakis (1994) via a simple
combination of randomized rounding
($\mathbf{Pr}\left[x^{\text{int}}_{i}=1\right]=x_{i}^{*}$) and deterministic
rounding ($x^{\text{int}}_{i}=1$ if $x^{*}_{i}\geq p$, where commonly
$p=0.5$).

We attempt to round the fractional solution $R^{*},S^{*}$ using these two
strategies, and then apply Algorithm 1 to $R^{\text{int}},S^{\text{int}}$.
However, direct application of deterministic rounding returns infeasible
results: the rounded solution violates constraints. Randomized rounding may
show more promise as a single relaxed solution can be used to sample many
integral solutions, some of which are hopefully feasible. Unfortunately, using
randomized rounding with the LP relaxation for VGG16 at a 4$\times$ smaller
budget than default, we could not find a single feasible solution out of
50,000 samples.

Algorithm 2 Two-phase rounding

Input: Fractional checkpoint matrix $S^{*}$ from LP

Output: Binary $S^{\text{int}}$, $R^{\text{int}}$, Free

Round $S^{*}$ deterministically:
$S_{t,i}^{\text{int}}\leftarrow\mathbbm{1}[S_{t,i}^{*}>0.5]$

$R^{\text{int}}\leftarrow\mathbf{I}_{n}$ thereby satisfying (7a)

while $\exists t\geq 2,i\in[n]$ such that
$S^{\text{int}}_{t,i}>R^{\text{int}}_{t-1,i}+S^{\text{int}}_{t-1,i}$ i.e.
(4.3) violated do

Compute $v_{i}$ to materialize checkpoint: $R^{\text{int}}_{t-1,i}\leftarrow
1$

end while

while $\exists t\geq 1,(i,j)\in E$ such that
$R^{\text{int}}_{t,j}>R^{\text{int}}_{t,i}+S^{\text{int}}_{t,i}$i.e. (4.3)
violated do

Compute $v_{i}$ as temporary for dependency: $R^{\text{int}}_{t,i}\leftarrow
1$

end while

Evaluate Free by simulating execution

return $S^{\text{int}}$, $R^{\text{int}}$, Free

###  5.2 A two-phase rounding strategy

To find feasible solutions, we introduce two-phase rounding, detailed in
Algorithm 2. Two-phase rounding is applicable when a subset of variables can
be solved in polynomial time given the remaining variables. Our approximation
algorithm only rounds the checkpoint matrix $S^{*}$. Given $S^{*}$, we solve
for the conditionally optimal binary computation matrix $R^{\text{int}}$ by
setting as few values to $1$ as possible. Algorithm 2 begins with an all-zero
matrix $R^{\text{int}}=0$, then iteratively corrects violated correctness
constraints.

Note that during any of the above steps, once we set some
$R_{i,j}^{\text{int}}=1$, the variable is never changed. Algorithm 2 corrects
constraints in a particular order so that constraints that are satisfied will
continue to be satisfied as other violated constraints are corrected. The
matrix $R^{\text{int}}$ generated by this rounding scheme will be optimal up
to the choice of $S^{\text{int}}$ as every entry in $R^{\text{int}}$ is set to
$1$ if and only if it is necessary to satisfy a constraint. In implementation,
we detect and correct violations of (4.3) in reverse topological order for
each stage, scanning $R^{\text{int}},S^{\text{int}}$ matrices from right to
left.

###  5.3 Memory budget feasibility

Since we approximate $S$ by rounding the fractional solution,
$S^{\text{int}},R^{\text{int}}$ can be infeasible by the budget constraint
$U_{t,k}\leq M_{\text{budget}}$. While the fractional solution may come under
the budget and two-phase rounding preserves correctness constraints, the
rounding procedure makes no attempt to maintain budget feasibility. Therefore,
we leave an allowance on the total memory budget constraint
($U_{t,k}\leq(1-\epsilon)M_{\text{budget}}$). We empirically find
$\epsilon=0.1$ to work well.

Figure 5: Computational overhead versus memory budget for (a) VGG16 image
classification NN Simonyan & Zisserman (2014), (b) MobileNet image
classification NN, and (c) the U-Net semantic segmentation NN Ronneberger et
al. (2015). Overhead is with respect to the best possible strategy without a
memory restriction based on a profile-based cost model of a single NVIDIA V100
GPU. For U-Net (c), at the 16 GB V100 memory budget, we achieve a $1.20\times$
speedup over the best baseline—linearized greedy—and a $1.38\times$ speedup
over the next best—linearized $\sqrt{n}$. Takeaway: our model- and hardware-
aware solver produces in-budget solutions with the lowest overhead on linear
networks (a-b), and dramatically lowers memory consumption and overhead on
complex architectures (c).

##  6 Evaluation

In this section, we investigate the impact of tensor rematerialization on the
cost and memory usage of DNN training. We study the following experimental
questions: (1) What is the trade-off between memory usage and computational
overhead when using rematerialization? (2) Are large inputs practical with
rematerialization? and (3) How well can we approximate the optimal
rematerialization policy?

We compare our proposed solver against baseline heuristics on representative
image classification and high resolution semantic segmentation models
including VGG16, VGG19, ResNet50, MobileNet, U-Net and FCN with VGG layers,
and SegNet. As prior work is largely limited to linear graphs, we propose
novel extensions where necessary for comparison. Results show that optimal
rematerialization allows significantly lower computational overhead than
baselines at all memory budgets, and lower memory usage than previously
possible. As a consequence, optimal rematerialization allows training with
larger input sizes than previously possible, up to 5.1$\times$ higher batch
sizes on the same accelerator. Finally, we find that our two-phase rounding
approximation algorithm finds near-optimal solutions in polynomial time.

###  6.1 Baselines and generalizations

Table 1 summarizes baseline rematerialization strategies. The nominal
evaluation strategy stores all features generated during the forward pass for
use during the backward pass—this is the default in frameworks such as
TensorFlow. Hence, every layer is computed once. We refer to this baseline as
Checkpoint all, an ideal approach given ample memory.

On the linear graph architectures, such as VGG16 and MobileNet (v1), we
directly apply prior work from Griewank & Walther (2000) and Chen et al.
(2016b), baselines referred to as Griewank and Walther $\log{n}$, Chen et al.
$\sqrt{n}$ and Chen et al. greedy. To build a tradeoff curve for computation
versus memory budget, we search over the segment size hyperparameter $b$ in
the greedy strategy. However, these baselines cannot be used for modern
architectures with residual connections. For a fair comparison, we extend the
$\sqrt{n}$ and greedy algorithms to apply to general computation graphs (e.g.
ResNet50 and U-Net).

Chen et al. (2016b) suggests manually annotating good checkpointing candidates
in a computation graph. For the first extensions, denoted by AP $\sqrt{n}$ and
AP greedy, we automatically identify articulation points, or cut vertices,
vertices that disconnect the forward pass DAG, and use these as candidates.
The heuristics then select a subset of these candidates, and we work backwards
from the checkpoints to identify which nodes require recomputation.

Still, some networks have few articulation points, including U-Net. We also
extend heuristics by treating the original graph as a linear network, with
nodes connected in topological order, again backing out the minimal
recomputations from the selected checkpoints. These extensions are referred to
as Linearized $\sqrt{n}$ and Linearized greedy.

Sections B.1 and B.2 provide more details on our generalizations. Note that
all proposed generalizations exactly reproduce the original heuristics on
linear networks.

###  6.2 Evaluation setup

Checkmate is implemented in Tensorflow 2.0 Abadi et al. (2016), accepting
user-defined models expressed via the high-level Keras interface. We extract
the forward and backward computation graph, then construct and solve
optimization problem (4.7) with the Gurobi mathematical programming library as
an integer linear program. Finally, Checkmate translates solutions into
execution plans and constructs a new static training graph. Together, these
components form the Checkmate system, illustrated in Figure 2.

To accelerate problem construction, decision variables $R$ and $S$ are
expressed as lower triangular matrices, as are accounting variables $U$. Free
is represented as a $|V|\times{}|E|$ matrix. Except for our maximum batch size
experiments, solutions are generated with a user-configurable time limit of
$3600$ seconds, though the majority of problems solve within minutes. Problems
with exceptionally large batch sizes or heavily constrained memory budgets may
reach this time limit while the solver attempts to prove that the problem is
infeasible. The cost of a solution is measured with a profile-based cost model
(Section 4.10) and compared to the ideal, unachievable cost with no
recomputation.

The feasible set of our optimal ILP formulation is a superset of baseline
heuristics. We implement baselines as a static policy for the decision
variable $S$ and then solve for the lowest-cost recomputation schedule using a
similar procedure to that described in Algorithm 2.

###  6.3 What is the trade-off between memory usage and computational
overhead?

Figure 5 compares remateralization strategies on VGG-16, MobileNet, and U-Net.
The y-axis shows the computational overhead of checkpointing in terms of time
as compared to baseline. The time is computed by profiling each individual
layer of the network. The x-axis shows the total memory budget required to run
each model with the specified batch size, computed for single precision
training. Except for the $\sqrt{n}$ heuristics, each rematerialization
algorithm has a knob to trade-off the amount of recomputation and memory
usage, where a smaller memory budget leads to higher overhead.

Takeaways: For all three DNNs, Checkmate produces clearly faster execution
plans as compared to algorithms proposed by Chen et al. (2016b) and Griewank &
Walther (2000) – over $1.2\times$ faster than the next best on U-Net at the
NVIDIA V100 memory budget. Our framework allows training a U-Net at a batch
size of 32 images per GPU with less than 10% higher overhead. This would
require 23 GB of memory without rematerialization, or with the original
baselines without our generalizations.

###  6.4 Are large inputs practical with rematerialization?

Figure 6:  Maximum batch size possible on a single NVIDIA V100 GPU when using
different generalized rematerialization strategies with at most a single extra
forward pass. We enable increasing batch size by up to _$5.1\times$_ over the
current practice of caching all activations (on MobileNet), and up to
$1.73\times$ over the best checkpointing scheme (on U-Net).

The maximum batch size enabled by different rematerialization strategies is
shown in Figure 6. The y-axis shows the theoretical maximum batch size we
could feasibly train with bounded compute cost. This is calculated by
enforcing that the total cost must be less than the cost of performing just
one additional forward pass. That is, in Figure 6 the cost is at most an
additional forward pass higher, _if_ the specified batch size would have fit
in GPU memory. To find Checkmate’s maximum batch size, we reformulate Problem
(4.7) to maximize a batch size variable $B\in\mathbb{N}$ subject to modified
memory constraints that use $B*M_{i}$ in place of $M_{i}$ and subject to an
additional cost constraint,

$\sum_{t=1}^{n}\sum_{i=1}^{t}C_{i}R_{t,i}\leq 2\sum_{v_{i}\in
G_{\text{fwd}}}C_{i}+\sum_{v_{i}\in G_{\text{bwd}}}C_{i}.$ (8)

The modified integer program has quadratic constraints, and is difficult to
solve. We set a time limit of one day for the experiment, but Gurobi may be
unable to reach optimality within that limit. Figure 6 then provides a lower
bound on the maximum batch size that Checkmate can achieve.

For fair comparison on the non-linear graphs used in U-Net, FCN, and ResNet,
we use the AP $\sqrt{n}$ and linearized greedy baseline generalizations
described in Section 6.1. For the baselines, we iterate over batch sizes, find
candidate solutions (multiple candidates for linearized greedy), and filter
out the solutions that cost more than an additional forward pass or that would
exceed the 16GB memory budget. The iteration stops when no solutions are
available.

Costs are measured in FLOPs, determined statically. U-Net, FCN8 and SegNet
semantic segmentation networks use a resolution of $416\times 608$, and
classification networks ResNet50, VGG19 and MobileNet use resolution
$224\times 224$.

Takeaways: We can theoretically increase the batch size of U-Net to $61$ at a
high resolution, an unprecedented result. For many tasks such as semantic
segmentation, where U-Net is commonly used, it is not possible to use batch
sizes greater than $16$, depending on resolution. This is sub-optimal for
batch normalization layers, and being able to increase the batch size by
$3.8\times$ ($61$ vs $16$ at this resolution) is quite significant. Orthogonal
approaches to achieve this include model parallelism and distributed memory
batch normalization which can be significantly more difficult to implement and
have high communication costs.

Furthermore, for MobileNet, Checkmate allows a batch size of $1105$ which is
$1.73\times$ higher than the best baseline solution, a greedy heuristic, and
$5.1\times$ common practice, checkpointing all activations. The same schedules
can also be used to increase image resolution rather than batch size.

###  6.5 How well can we approximate the optimal rematerialization policy?

To understand how well our LP rounding strategy (Section 5) approximates the
ILP, we measure the ratio
$\textsc{Cost}_{\text{approx}}/\textsc{Cost}_{\text{opt}}$, i.e. the speedup
of the optimal schedule, in FLOPs. As in Section 6.3, we solve each strategy
at a range of memory budgets, then compute the geometric mean of the ratio
across budgets. The aggregated ratio is used because some budgets are feasible
via the ILP but not via the approximations. Table 2 shows results. The two-
phase deterministic rounding approach has approximation factors close to
optimal, at most $1.06\times$ for all tested architectures.

|  | AP
---
$\sqrt{n}$
| AP
---
greedy
| Griewank
---
$\log n$
| Two-phase
---
LP rounding
MobileNet | 1.14$\times$ | 1.07$\times$ | 7.07$\times$ | 1.06$\times$
VGG16 | 1.28$\times$ | 1.06$\times$ | 1.44$\times$ | 1.01$\times$
VGG19 | 1.54$\times$ | 1.39$\times$ | 1.75$\times$ | 1.00$\times$
U-Net | 1.27$\times$ | 1.23$\times$ | - | 1.03$\times$
ResNet50 | 1.20$\times$ | 1.25$\times$ | - | 1.05$\times$

Table 2: Approximation ratios for baseline heuristics and our LP rounding
strategy. Results are given as the geometric mean speedup of the optimal ILP
across feasible budgets.

##  7 Conclusion

One of the main challenges when training large neural networks is the limited
capacity of high-bandwidth memory on accelerators such as GPUs and TPUs. This
has created a memory wall that limits the size of the models that can be
trained. The bottleneck for state-of-the-art model development is now memory
rather than data and compute availability, and we expect this trend to worsen
in the future.

To address this challenge, we proposed a novel rematerialization algorithm
which allows large models to be trained with limited available memory. Our
method does not make the strong assumptions required in prior work, supporting
general non-linear computation graphs such as residual networks and capturing
the impact of non-uniform memory usage and computation cost throughout the
graph with a hardware-aware, profile-guided cost model. We presented an ILP
formulation for the problem, implemented the Checkmate system for optimal
rematerialization in TensorFlow, and tested the proposed system on a range of
neural network models. In evaluation, we find that optimal rematerialization
has minimal computational overhead at a wide range of memory budgets and
showed that Checkmate enables practitioners to train high-resolution models
with significantly larger batch sizes. Finally, a novel two-phase rounding
strategy closely approximates the optimal solver.

## Acknowledgements

We would like to thank Barna Saha and Laurent El Ghaoui for guidance on
approximation, Mong H. Ng for help in evaluation, and the paper and artifact
reviewers for helpful suggestions. In addition to NSF CISE Expeditions Award
CCF-1730628 and ONR PECASE N000141612723, this work was supported by gifts
from Alibaba, Amazon Web Services, Ant Financial, CapitalOne, Ericsson,
Facebook, Futurewei, Google, Intel, Microsoft, NVIDIA, Scotiabank, Splunk and
VMware. This work was also supported by the NSF GRFP under Grant No.
DGE-1752814. Any opinions, findings, and conclusions or recommendations
expressed in this material are those of the author(s) and do not necessarily
reflect the views of the NSF.

## References

  * Abadi et al. (2016) Abadi, M., Agarwal, A., Barham, P., Brevdo, E., Chen, Z., Citro, C., Corrado, G. S., Davis, A., Dean, J., Devin, M., Ghemawat, S., Goodfellow, I., Harp, A., Irving, G., Isard, M., Jia, Y., Jozefowicz, R., Kaiser, L., Kudlur, M., Levenberg, J., Mane, D., Monga, R., Moore, S., Murray, D., Olah, C., Schuster, M., Shlens, J., Steiner, B., Sutskever, I., Talwar, K., Tucker, P., Vanhoucke, V., Vasudevan, V., Viegas, F., Vinyals, O., Warden, P., Wattenberg, M., Wicke, M., Yu, Y., and Zheng, X.  TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems.  March 2016.
  * Beaumont et al. (2019) Beaumont, O., Herrmann, J., Pallez, G., and Shilova, A.  Optimal memory-aware backpropagation of deep join networks.  Research Report RR-9273, Inria, May 2019.
  * Briggs et al. (1992) Briggs, P., Cooper, K. D., and Torczon, L.  Rematerialization.  In _Proceedings of the ACM SIGPLAN 1992 Conference on Programming Language Design and Implementation_ , PLDI ’92, pp. 311–321, New York, NY, USA, 1992.
  * Brock et al. (2018) Brock, A., Donahue, J., and Simonyan, K.  Large scale GAN training for high fidelity natural image synthesis.  _arXiv preprint arXiv:1809.11096_ , 2018.
  * Bulo et al. (2018) Bulo, S. R., Porzi, L., and Kontschieder, P.  In-place Activated BatchNorm for Memory-Optimized Training of DNNs.  In _2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 5639–5647. IEEE, June 2018.
  * Canziani et al. (2016) Canziani, A., Paszke, A., and Culurciello, E.  An Analysis of Deep Neural Network Models for Practical Applications.  May 2016.  arXiv: 1605.07678.
  * Chaitin et al. (1981) Chaitin, G. J., Auslander, M. A., Chandra, A. K., Cocke, J., Hopkins, M. E., and Markstein, P. W.  Register allocation via coloring.  _Computer Languages_ , 6(1):47–57, January 1981\.
  * Chen et al. (2016a) Chen, L.-C., Papandreou, G., Kokkinos, I., Murphy, K., and Yuille, A. L.  DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs.  June 2016a.  arXiv: 1606.00915.
  * Chen et al. (2016b) Chen, T., Xu, B., Zhang, C., and Guestrin, C.  Training Deep Nets with Sublinear Memory Cost.  April 2016b.  arXiv: 1604.06174.
  * Chen et al. (2017) Chen, X., Ma, H., Wan, J., Li, B., and Xia, T.  Multi-view 3D Object Detection Network for Autonomous Driving.  In _2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 6526–6534. IEEE, 2017.
  * Child et al. (2019) Child, R., Gray, S., Radford, A., and Sutskever, I.  Generating Long Sequences with Sparse Transformers.  April 2019.  arXiv: 1904.10509.
  * Cytron et al. (1991) Cytron, R., Ferrante, J., Rosen, B. K., Wegman, M. N., and Zadeck, F. K.  Efficiently Computing Static Single Assignment Form and the Control Dependence Graph.  _ACM Trans. Program. Lang. Syst._ , 13(4):451–490, October 1991.
  * Dai et al. (2019) Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q. V., and Salakhutdinov, R.  Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context.  January 2019.  arXiv: 1901.02860.
  * Devlin et al. (2018) Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K.  BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.  October 2018.  arXiv: 1810.04805.
  * Dong et al. (2016) Dong, C., Loy, C. C., He, K., and Tang, X.  Image super-resolution using deep convolutional networks.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 38(2):295–307, Feb 2016.
  * Feng & Huang (2018) Feng, J. and Huang, D.  Cutting Down Training Memory by Re-fowarding.  July 2018.
  * Forrest et al. (2019) Forrest, J. J., Vigerske, S., Ralphs, T., Santos, H. G., Hafer, L., Kristjansson, B., Fasano, J., Straver, E., Lubin, M., rlougee, jpgoncal1, Gassmann, H. I., and Saltzman, M.  COIN-OR Branch-and-Cut solver, June 2019.
  * Gholami et al. (2018) Gholami, A., Azad, A., Jin, P., Keutzer, K., and Buluc, A.  Integrated model, batch, and domain parallelism in training neural networks.  In _Proceedings of the 30th on Symposium on Parallelism in Algorithms and Architectures_ , pp. 77–86. ACM, 2018.
  * (19) GLPK.  GNU Project - Free Software Foundation (FSF).
  * Gomez et al. (2017) Gomez, A. N., Ren, M., Urtasun, R., and Grosse, R. B.  The Reversible Residual Network: Backpropagation Without Storing Activations.  In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 30_ , pp. 2214–2224. Curran Associates, Inc., 2017.
  * Goodwin & Wilken (1996) Goodwin, D. W. and Wilken, K. D.  Optimal and Near-optimal Global Register Allocation Using 0–1 Integer Programming.  _Software: Practice and Experience_ , 26(8):929–965, 1996.
  * Griewank & Walther (2000) Griewank, A. and Walther, A.  Algorithm 799: revolve: an implementation of checkpointing for the reverse or adjoint mode of computational differentiation.  _ACM Transactions on Mathematical Software_ , 26(1):19–45, March 2000.
  * Gruslys et al. (2016) Gruslys, A., Munos, R., Danihelka, I., Lanctot, M., and Graves, A.  Memory-efficient Backpropagation Through Time.  In _Proceedings of the 30th International Conference on Neural Information Processing Systems_ , NIPS’16, pp. 4132–4140, USA, June 2016. Curran Associates Inc.
  * Gueguen et al. (2018) Gueguen, L., Sergeev, A., Kadlec, B., Liu, R., and Yosinski, J.  Faster Neural Networks Straight from JPEG.  In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 31_ , pp. 3933–3944. Curran Associates, Inc., 2018.
  * He et al. (2016) He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 770–778, 2016.
  * Holder (2008) Holder, L.  Graph Algorithms: Applications, 2008.
  * Huang et al. (2017) Huang, G., Liu, Z., Van Der Maaten, L., and Weinberger, K. Q.  Densely connected convolutional networks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 4700–4708, 2017.
  * Ioffe & Szegedy (2015) Ioffe, S. and Szegedy, C.  Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift.  _International Conference on Machine Learning_ , February 2015.
  * Jain et al. (2018) Jain, A., Phanishayee, A., Mars, J., Tang, L., and Pekhimenko, G.  Gist: Efficient Data Encoding for Deep Neural Network Training.  In _Proceedings of the 45th Annual International Symposium on Computer Architecture_ , ISCA ’18, pp. 776–789, Piscataway, NJ, USA, 2018. IEEE Press.
  * Jain et al. (2019) Jain, P., Mo, X., Jain, A., Tumanov, A., Gonzalez, J. E., and Stoica, I.  The ooo vliw jit compiler for gpu inference.  _arXiv preprint arXiv:1901.10008_ , 2019.
  * Jia et al. (2018a) Jia, Z., Lin, S., Qi, C. R., and Aiken, A.  Exploring Hidden Dimensions in Accelerating Convolutional Neural Networks.  In _International Conference on Machine Learning_ , pp. 2274–2283, July 2018a.
  * Jia et al. (2018b) Jia, Z., Zaharia, M., and Aiken, A.  Beyond Data and Model Parallelism for Deep Neural Networks.  _SysML Conference_ , pp. 13, Feb. 2018b.
  * Karmarkar (1984) Karmarkar, N.  A new polynomial-time algorithm for linear programming.  In _Proceedings of the sixteenth annual ACM symposium on Theory of computing_ , pp. 302–311. ACM, 1984.
  * Kim et al. (2016) Kim, J., Lee, J. K., and Lee, K. M.  Accurate image super-resolution using very deep convolutional networks.  In _2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 1646–1654, June 2016.  doi: 10.1109/CVPR.2016.182.
  * Koes & Goldstein (2006) Koes, D. R. and Goldstein, S. C.  A Global Progressive Register Allocator.  In _Proceedings of the 27th ACM SIGPLAN Conference on Programming Language Design and Implementation_ , PLDI ’06, pp. 204–215, New York, NY, USA, 2006. ACM.  event-place: Ottawa, Ontario, Canada.
  * Krizhevsky et al. (2012) Krizhevsky, A., Sutskever, I., and Hinton, G. E.  ImageNet Classification with Deep Convolutional Neural Networks.  In Pereira, F., Burges, C. J. C., Bottou, L., and Weinberger, K. Q. (eds.), _Advances in Neural Information Processing Systems 25_ , pp. 1097–1105. Curran Associates, Inc., 2012.
  * Lattner (2002) Lattner, C.  LLVM: An Infrastructure for Multi-Stage Optimization.  Master’s thesis, Computer Science Dept., University of Illinois at Urbana-Champaign, Urbana, IL, December 2002.
  * Liu et al. (2019) Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V.  RoBERTa: A Robustly Optimized BERT Pretraining Approach.  July 2019.  arXiv: 1907.11692.
  * Long et al. (2015) Long, J., Shelhamer, E., and Darrell, T.  Fully convolutional networks for semantic segmentation.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 3431–3440, 2015.
  * Lozano et al. (2018) Lozano, R. C., Carlsson, M., Blindell, G. H., and Schulte, C.  Combinatorial Register Allocation and Instruction Scheduling.  April 2018.  arXiv: 1804.02452.
  * (41) McCandlish, S., Kaplan, J., Amodei, D., and Team, O. D.  An Empirical Model of Large-Batch Training.  arXiv: 1812.06162.
  * Meng et al. (2017) Meng, C., Sun, M., Yang, J., Qiu, M., and Gu, Y.  Training Deeper Models by GPU Memory Optimization on TensorFlow.  pp. 8, December 2017.
  * Micikevicius (2011) Micikevicius, P.  Local Memory and Register Spilling, 2011.
  * Nakata (1967) Nakata, I.  On Compiling Algorithms for Arithmetic Expressions.  _Commun. ACM_ , 10(8):492–494, August 1967.  ISSN 0001-0782.  doi: 10.1145/363534.363549.
  * Nesterov & Nemirovskii (1994) Nesterov, Y. and Nemirovskii, A.  _Interior-point polynomial algorithms in convex programming_ , volume 13.  Siam, 1994.
  * NVIDIA (2017) NVIDIA.  NVIDIA Tesla V100 GPU Architecture, 2017.
  * Olesen (2011) Olesen, J. S.  Register Allocation in LLVM 3.0, November 2011.
  * Paszke et al. (2017) Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., Lin, Z., Desmaison, A., Antiga, L., and Lerer, A.  Automatic differentiation in PyTorch.  In _NIPS 2017 Autodiff Workshop_ , 2017.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S.  Pytorch: An imperative style, high-performance deep learning library.  In Wallach, H., Larochelle, H., Beygelzimer, A., d’ Alché-Buc, F., Fox, E., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 32_ , pp. 8024–8035. Curran Associates, Inc., 2019.
  * Pohlen et al. (2017) Pohlen, T., Hermans, A., Mathias, M., and Leibe, B.  Full-resolution residual networks for semantic segmentation in street scenes.  In _Computer Vision and Pattern Recognition (CVPR), 2017 IEEE Conference on_ , 2017.
  * Punjani (2004) Punjani, M.  Register Rematerialization in GCC.  In _GCC Developers’ Summit_ , volume 2004. Citeseer, 2004\.
  * Ronneberger et al. (2015) Ronneberger, O., Fischer, P., and Brox, T.  U-Net: Convolutional Networks for Biomedical Image Segmentation.  In Navab, N., Hornegger, J., Wells, W. M., and Frangi, A. F. (eds.), _Medical Image Computing and Computer-Assisted Intervention – MICCAI 2015_ , Lecture Notes in Computer Science, pp. 234–241. Springer International Publishing, 2015.  ISBN 978-3-319-24574-4.
  * Rosen et al. (1988) Rosen, B. K., Wegman, M. N., and Zadeck, F. K.  Global Value Numbers and Redundant Computations.  In _Proceedings of the 15th ACM SIGPLAN-SIGACT Symposium on Principles of Programming Languages_ , POPL ’88, pp. 12–27, New York, NY, USA, 1988. ACM.
  * Sethi (1973) Sethi, R.  Complete Register Allocation Problems.  pp. 14, April 1973.
  * Simonyan & Zisserman (2014) Simonyan, K. and Zisserman, A.  Very Deep Convolutional Networks for Large-Scale Image Recognition.  September 2014.  arXiv: 1409.1556.
  * Siskind & Pearlmutter (2018a) Siskind, J. M. and Pearlmutter, B. A.  Divide-and-conquer checkpointing for arbitrary programs with no user annotation.  _Optimization Methods and Software_ , 33(4-6):1288–1330, 2018a.  doi: 10.1080/10556788.2018.1459621.
  * Siskind & Pearlmutter (2018b) Siskind, J. M. and Pearlmutter, B. A.  Divide-and-Conquer Checkpointing for Arbitrary Programs with No User Annotation.  _Optimization Methods and Software_ , 33(4-6):1288–1330, November 2018b.
  * Sivathanu et al. (2019) Sivathanu, M., Chugh, T., Singapuram, S. S., and Zhou, L.  Astra: Exploiting Predictability to Optimize Deep Learning.  In _Proceedings of the Twenty-Fourth International Conference on Architectural Support for Programming Languages and Operating Systems - ASPLOS ’19_ , pp. 909–923, Providence, RI, USA, 2019\. ACM Press.
  * Sze et al. (2017) Sze, V., Chen, Y.-H., Yang, T.-J., and Emer, J. S.  Efficient processing of deep neural networks: A tutorial and survey.  _Proceedings of the IEEE_ , 105(12):2295–2329, 2017.
  * Szegedy et al. (2015) Szegedy, C., Wei Liu, Yangqing Jia, Sermanet, P., Reed, S., Anguelov, D., Erhan, D., Vanhoucke, V., and Rabinovich, A.  Going deeper with convolutions.  In _2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 1–9, June 2015.  doi: 10.1109/CVPR.2015.7298594.
  * Tai et al. (2017) Tai, Y., Yang, J., and Liu, X.  Image super-resolution via deep recursive residual network.  In _2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 2790–2798, July 2017.  doi: 10.1109/CVPR.2017.298.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I.  Attention is All you Need.  In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 30_ , pp. 5998–6008. Curran Associates, Inc., 2017.
  * Wu & He (2018) Wu, Y. and He, K.  Group Normalization.  pp. 3–19, 2018.
  * Xie et al. (2017) Xie, S., Girshick, R., Dollár, P., Tu, Z., and He, K.  Aggregated residual transformations for deep neural networks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 1492–1500, 2017.
  * Yang et al. (2018) Yang, B., Liang, M., and Urtasun, R.  HDNET: Exploiting HD Maps for 3D Object Detection.  pp. 10, 2018.
  * Yannakakis (1994) Yannakakis, M.  On the approximation of maximum satisfiability.  _Journal of Algorithms_ , 17(3):475–502, 1994\.

Figure 7: We visualize schedules ($R$ matrix) for VGG19. The $R$ matrix
denotes when each layer in a neural network is evaluated. Rows refer to
execution stage while columns refer to a particular layer. We make no
distinction between forward and backward nodes. TensorFlow 2.0 can train VGG19
on a V100 with a batch size of 167. By applying the heuristic from Chen et al.
(2016b), a V100 can sustain training at a batch size od 197. However, our
proposed method in Checkmate sustains training at a batch size of 289,
representing a 73% improvement. The integer linear program took 10 seconds to
solve to optimality.

##  Appendix A Integrality gap

To understand why the partitioned variant of the MILP (Section 4.2) is faster
to solve via branch-and-bound, we can measure the integrality gap for
particular problem instances. The integrality gap is the maximum ratio between
the optimal value of the ILP and its relaxation, defined as follows:

$\displaystyle IG=\max_{I}\frac{\textsc{Cost}_{int}}{\textsc{Cost}_{frac}},$

where $\textsc{Cost}_{int}$ and $\textsc{Cost}_{frac}$ are the optimal value
the ILP and that of its relaxation, respectively. $I$ describes a problem
instance,

$I=(G,C,M,M_{\text{budget}}).$

As our ILP is a minimization problem,
$\textsc{Cost}_{int}\geq\textsc{Cost}_{frac}$ for all $I$, and $IG\geq 1$.
While it is not possible to measure the ratio between the ILP and LP solutions
for all problem instances, the ratio for any particular problem instance gives
a lower bound on the integrality gap.

For the 8-layer linear neural network graph discussed in Section 4.2,
frontier-advancement reduces the integrality gap from 21.56 to 1.18, i.e. the
LP relaxation is significantly tighter. In branch-and-bound algorithms for ILP
optimiztion, a subset of feasible solutions can be pruned if the LP relaxation
over the subset yields an objective higher than the best integer solution
found thus far. With a tight LP relaxation, this condition for pruning is
often met, so fewer solutions need to be enumerated.

##  Appendix B Generalizations of prior work

###  B.1 AP $\sqrt{n}$ and AP greedy

We identify Articulation Points (AP) in the undirected form of the forward
pass data-flow graph as candidates for checkpointing. APs are vertices that
increase the number of connected components (i.e. disconnect) the graph if
removed, and can be identified in time $O(V+E)$ via a modified DFS traversal
Holder (2008). An articulation point $v_{a}$ is a good candidate for
checkpointing as subsequent vertices in the topological order have no
dependencies on vertices before $v_{a}$ in the order. DNN computation graphs
are connected, so each intermediate tensor can be reconstructed from a single
articulation point earlier in the topological order, or the input if there is
no such AP. APs include the input and output nodes of residual blocks in
ResNet, but not vertices inside blocks. We apply Chen’s heuristics to
checkpoint a subset of these candidates, then solve for the optimal
recomputation plan $R$ to restore correctness. Solving for $R$ ensures that a
node’s dependencies are resident prior to evaluation.

We could find $R$ by solving the optimization problem (4.7) with additional
constraints on $S$ that encode the heuristically selected checkpoints.
However, as $S$ is given, the optimization is solvable in $O(|V||E|)$ via a
graph traversal per row of $R$ that fills in entries when a needed value is
not in memory by the same process described in Section 5.2.

###  B.2 Linearized $\sqrt{n}$ and Linearized greedy

The forward graph of the DNN
${G_{\text{fwd}}=(V_{\text{fwd}},E_{\text{fwd}})}$ can be treated as a linear
graph ${G_{\text{lin}}=(V_{\text{fwd}},E_{\text{lin}})}$ with edges connecting
consecutive vertices in a topological order:

$E_{\text{lin}}=\\{(v_{1},v_{2}),(v_{2},v_{3}),\ldots,(v_{L-1},v_{L})\\}$

While $G_{\text{lin}}$ does not properly encode data dependencies, it is a
linear graph that baselines can analyze. To extend a baseline, we apply it to
$G_{\text{lin}}$, generate checkpoint matrix $S$ from the resulting checkpoint
set, and find the optimal $R$ as with the AP baselines.

##  Appendix C Hardness of rematerialization

Sethi (1973) reduced 3-SAT to a decision problem based on register allocation
in straight line programs, with no recomputation permitted. Such programs can
be represented by result-rooted Directed Ayclic Graphs (DAGs), with nodes
corresponding to operations and edges labeled by values. In Sethi’s graphs,
the desired results are the roots of the DAG. If a program has no common
subexpressions, i.e. the graph forms a tree, optimal allocation is possible
via a linear time tree traversal Nakata (1967). However, Sethi’s reduction
shows a register allocation decision problem in the general case—whether a
result-rooted DAG can be computed with fewer than $k$ registers without
recomputation—is NP-complete.

The decision problem characterizes computation of a DAG as a sequence of four
possible moves of stones, or registers, on the nodes of the graph, analogous
to statements discussed in Section 4.9. The valid moves are to (1) place a
register at a leaf, computing it, or (2) pick up a register from a node. Also,
if there are registers at all children of a node $x$, then it is valid to (3)
place a register at $x$, computing it, or (4) move a stone to $x$ from one of
the children of $x$, computing $x$. The register allocation problem reduces to
the following no-overhead rematerialization decision problem (RP-DEC):

######  Definition C.1.

(RP-DEC): Given result-terminated data-flow DAG $G=(V,E)$ corresponding to a
program, with unit cost to compute each node and unit memory for the results
of each node, does there exist an execution plan that evaluates the leaf
(terminal) node $t\in V$ with maximum memory usage $b$ at cost at most $|V|$?

RP-DEC is decidable by solving the memory-constrained form of Problem 1 with
sufficient stages, then checking if the returned execution plan has cost at
most $|V|$. RP-DEC closely resembles Sethi’s decision problem, differing only
in subtleties. The register allocation DAG is rooted at the desired result $t$
whereas a data-flow graph terminates at the result. Second, register-based
computations can be in place, e.g. a summation $a+b$ may be written to the
same location as either of the operands. In neural network computation graphs,
we cannot perform all computations in place, so we did not make this
assumption. To reduce Sethi’s decision problem to RP-DEC, given result-rooted
DAG $G$, construct result-terminated $G^{\prime}$ by reversing all edges.
Then, if Sethi’s instance allows for at most $k$ registers, allow for a memory
budget of $b=k+1$ bytes: one byte to temporarily write outputs of operations
that would have been written in place.

Despite hardness of register allocation, Goodwin & Wilken (1996) observe that
a 0-1 integer program for optimal allocation under an instruction schedule has
empirical complexity $O(n^{2.5})$, polynomial in the number of constraints.
Similarly, Section 6 shows that the frontier-advancing, constrained
optimization problem (4.7) is tractable for many networks.

##  Appendix D Comparison of approximations

In Section 5, we discussed an approximation strategy based on rounding the LP
relaxation, evaluated with deterministic rounding in Section 6.5. Figure 8
compares schedules produced by our proposed two-phase rounding strategy when
the $S^{*}$ matrix from the LP relaxation is rounded with a randomized and a
deterministic approach. While two-phase randomized rounding of $S^{*}$ offers
a range of feasible solutions, two-phase deterministic rounding produces
consistently lower cost schedules. While appropriate for VGG16, for MobileNet,
our budget allowance $\epsilon=0.1$ is overly conservative as schedules use
less memory than the 16 GB budget. A search procedure over $\epsilon\in[0,1]$
could be used to produce more efficient schedules.

Figure 8: Comparison of the two-phase LP rounding approximation with
randomized rounding of $S^{*}$ and deterministic rounding of $S^{*}$ on
different models. We compare memory usage and computational cost (objective),
in milliseconds according to profile-based cost model. The average of the
randomized rounding costs is shown as a dotted line.


$\mathsf{CH}(k,b,\ell,s)$ with a perfect construction algorithm has
construction failure probability at most $\epsilon$.

###### Proof.

The chosen parameters incur an insertion failure if and only if any two items
hashed into the same entry. This probability is upper bounded by

${n\choose 2}\cdot\frac{1}{b}\leq n^{2}/b\leq\epsilon\implies
b=O(n^{2}/\epsilon)$

to complete the proof. ∎

###  B.2 Robustness

Using a large number of entries turns out to be a straightforward case to
enable robustness. We simply add it for completeness as it will require super-
polynomial in $n$ storage that is not feasible in most practical settings. The
result below is tight as it matches the non-robust lower bounds that we show
later.

######  Theorem 12.

For security parameter $\lambda$ and error $\epsilon=\mathsf{negl}(\lambda)$,
let $k=1$, $s=0$, $\ell=1$ and $b=O(1/\epsilon)$, then the cuckoo hashing
scheme $\mathsf{CH}(k,b,\ell,s)$ with a perfect construction algorithm is
$(\lambda,\epsilon)$-strongly robust.

###### Proof.

Let $U$ be the set of all items sent to the hash function where $|U|\leq
f(\lambda)$ for all $f(\lambda)=\mathsf{poly}(\lambda)$. Once again, we prove
that the probability that there exists any two items amongst the
$\mathsf{poly}(\lambda)$ items sent to the hash functions by the adversary are
mapped to the same entry is as follows:

$\Pr[\exists u_{1}\neq u_{2}\in U:H_{1}(u_{1})=H_{2}(u_{2})]\leq{|U|\choose
2}\cdot\frac{1}{b}\leq\frac{|U|^{2}}{b}.$

As $\epsilon=\mathsf{negl}(\lambda)$ and $|U|\leq\mathsf{poly}(\lambda)$, we
can set $b=O(1/\epsilon)$ to get that the above is smaller than $\epsilon$. ∎

###  B.3 Lower Bound

Finally, we present a lower bound that will match our two constructions above.

######  Theorem 13.

Let $k=O(1)$, $\ell=O(1)$, $b\geq n/\ell$ and $s=O(1)$. The failure
probability of $\mathsf{CH}(k,b,\ell,s)$ cuckoo hashing scheme where
$\mathcal{H}$ is a $(k\ell+s+1)$-wise independent hash function satisfies the
following:

$b=\Omega(1/\epsilon).$

###### Proof.

We can re-do the proof of Theorem 2 until we obtain the inequality

$k^{2}\ell+ks\geq\log(1/\epsilon)/\log(b/k).$

By plugging in that $k=O(1)$, $\ell=O(1)$ and $s=O(1)$, we obtain the desired
result of $b=\Omega(1/\epsilon)$. ∎

Note, if we choose typical parameters of $\epsilon=\mathsf{negl}(n)$, this
matches the $b=O(n^{2}/\epsilon)$ upper bounds of the prior constructions.

##  Appendix C Lower Bounds for Single Hash Function Settings

The lower bounds in Section 5.3 also immediately imply results in the single
hash function setting of $k=1$. However, it turns out stronger lower bounds
can be achieved using well-known results in the area of balls-and-bins. It has
already been proved that for $b=n$, the number of items assigned to a single
bin will be $\Omega(\log n/\log\log n)$ with probability except $1/n$ (for
example, see Lemma 5.12 in [MU17]). We re-prove the theorem for $b=\alpha\cdot
n$ for constant $\alpha>1$, but the proof techniques are identical. Note, this
justifies prior works usage of $k=2$ for parameter settings with large entries
$\ell$ [MP20] and stashes $s$ [KMW10].

######  Theorem 14.

Suppose that $k=1$ and $b=\alpha n$ for some constant $\alpha\geq 1$. If the
cuckoo hashing scheme $\mathsf{CH}(k,b,\ell,s)$ has failure probability at
most $\epsilon\leq 1/n$, then $\ell+s=\Omega(\log n/\log\log n)$.

###### Proof.

If we can prove that $t=\Omega(\log n/\log\log n)$ items are hashed into the
same entry with probability at least $\epsilon$, then we immediately complete
the proof as it must be that the single entry and stash must store all of the
$t$ allocated items implying that $\ell+s=\Omega(\log n/\log\log n)$.

We use the Poisson approximation of entry sizes where the number of items in
each entry is modelled using an independent Poisson variable with mean
$n/b=1/\alpha$. We denote the event $E$ when there exists one bin with at
least $t$ items. We note that $E$ is monotonically increasing in the number of
items $n$. Therefore, if we bound the probability of $E$ using the Poisson
approximation, we lose only a factor of $2$ when considering the true balls-
and-bins distribution using Corollary 5.11 in [MU17].

The probability that Poisson variable with mean $1/\alpha$ is greater than $t$
is at least $1/(\alpha^{t}\cdot t!\cdot e^{1/\alpha})$. The probability that
$n$ independent Poisson variables are all at most $t$ is

$\left(1-\frac{1}{\alpha^{t}\cdot t!\cdot e^{1/\alpha}}\right)^{n}\leq
e^{-n/(\alpha^{t}\cdot t!\cdot e^{1/\alpha})}.$

The above probability must be at most $\epsilon/2$ as we lose a factor of two
from the Poisson approximation. After re-arranging, it can be seen that

$\displaystyle e^{-n/(\alpha^{t}\cdot t!\cdot
e^{1/\alpha})}\leq\epsilon/2\implies t\log t=\Omega(\log
n+\log\log(1/\epsilon)).$

We note that $\log\log(1/\epsilon)=O(\log\log n)$, so we only need $t\log
t=\Omega(\log n)$. Therefore, setting $t=\Omega(\log n/\log\log n)$ suffices
to complete the proof. ∎

##  Appendix D Implications to Other Primitives

Recent works have leveraged techniques from cuckoo hashing to enable efficient
encodings of data such as cuckoo filters [FAKM14] and oblivious key-value
stores [GPR+21]. At a high level, both primitives take an input set
$\\{(\mathsf{id}_{1},v_{1}),\ldots,(\mathsf{id}_{n},v_{n})\\}$ of identifiers
and values. Furthermore, all $n$ values are typically random strings. Similar
to cuckoo hashing, each of the items are assigned to $k\geq 2$ entries
according to $k$ hash functions. For example, the $i$-th item
$(\mathsf{id}_{i},v_{i})$ is assigned to entries
$H_{1}(\mathsf{id}_{i}),\ldots,H_{k}(\mathsf{id}_{i})$. Then, the goal is to
find assignments to all entries such that the XOR of the $k$ entries will
equal $v_{i}$. If we let the resulting table be $T$, then we want to pick $T$
such that

$v_{i}=T[H_{1}(\mathsf{id}_{i})]\oplus\ldots\oplus T[H_{k}(\mathsf{id}_{i})]$

for all $i\in[n]$. For this setting, the failure probability corresponds to
the case that a table $T$ may not exist to satisfy the input set.

Lower Bounds. We can show that our lower bounds in Theorems 2 and 3 for fixed
choices of $\ell=1$ and $s=0$ also apply to these primitives. In the proof of
these theorems, the failure probability is lower bounded by computing the
probability that some subset of items $X$ end up being allocated to a subset
of entries $N(X)$ that is strictly smaller than $X$. That is, $|X|>|N(X)|$. In
this case, it is impossible to store $|X|$ items into strictly less than $|X|$
entries.

The same argument can also be applied to the above primitives. Here, there is
a system of $|X|$ linear equations with $|N(X)|<|X|$ variables. As the values
of the input set are random, it is impossible to solve the system of linear
equations. As a result, such a table $T$ will not exist. As a result, we can
see that the lower bounds in Theorems 2 and 3 with $\ell=1$ and $s=0$ also
apply to these primitives.

Upper Bounds. For constructions, we note that the same arguments can be
applied to show that if the underlying cuckoo hashing construction can
allocate the $n$ items, then there must exist a table $T$ whose entries
satisfy the constraints. The main challenge is to find efficient algorithms to
find such a table $T$. To our knowledge, the best algorithm would be using
Gaussian elimination. We leave it as an open problem to come up with
algorithms for computing the table $T$ using our constructions with large
numbers of hash functions.

Robustness with Private Keys. We also note that Filic et al. [FPUV22] studied
robust cuckoo filters in adversarial environments. However, this work
considered a slightly different model than our work. In particular, it was
shown that cuckoo filters may be constructed to be robust assuming that a
private key can be kept secret from the adversary.

##  Appendix E Cuckoo Hashing for Volume-Hiding Multi-Maps

Overview of [PPYY19]. To construct a volume-hiding multi-map, Patel et al.
[PPYY19] utilize cuckoo hashing in the following way. For pair of identifier
and value tuple, $(\mathsf{id},\mathbf{v}=(v_{1},\ldots,v_{k}))$, we convert
it into $k$ identifier and value tuples: $(\mathsf{id}\mid\mid
1,v_{1}),\ldots,(\mathsf{id}\mid\mid\ell,v_{k})$. We denote $k$ as the volume
of the identifier $\mathsf{id}$. After flattening all pairs of identifier and
value tuples into pairs of identifier and value, the resulting flattened pairs
are inserted into a cuckoo hashing table with $k=2$, $b=O(n)$, $\ell=1$ and
$s=O(1+\log(1/\epsilon)/\log n)$. To query for any identifier $\mathsf{id}$,
one simply performs cuckoo hashing queries for identifiers
$\mathsf{id}\mid\mid 1,\ldots,\mathsf{id}\mid\mid\ell$ where $z$ is the
maximum size (i.e., maximum volume) of any value tuple. There are two options
for storing the stash. In the original paper, the stash is stored by the
client and checked locally on each query. This version has $O(z)$ query
overhead but $O(1+\log(1/\epsilon)/\log n)$ client storage. If one wishes to
maintain $O(1)$ client storage, the stash may be encrypted and outsourced to
the server. Now, the client storage is $O(1)$ but the query overhead becomes
$O(z+\log(1/\epsilon)/\log n)$.

Modification using Improved Cuckoo Hashing. As a modification, we could
utilize our cuckoo hashing scheme with $k=O(1+\sqrt{\log(1/\epsilon)/\log
n})$, $b=O(n)$, $\ell=1$ and $s=0$. As there is no stash anymore, the client
storage remains $O(1)$. Instead, the query overhead now becomes
$O(z\cdot\sqrt{\log(1/\epsilon)/\log n})$. For small maximum volumes such as
$z=O(1)$, our construction has quadratically smaller overhead than the prior
construction in [PPYY19] with $O(1)$ client storage. As a special case, this
new construction would be useful when considering volume-hiding maps where the
goal is to simply hide whether an identifier exists in the map. For this
problem, the maximum volume is $z=1$.

##  Appendix F Running Time of Construction Algorithms

Throughout this section, we will rely heavily on the random bipartite graphs
that model cuckoo hashing that we described in Section 4.2. We point readers
back to that description for more details on the bipartite graph and the
importance of perfect left matchings.

We will consider several different construction algorithms for constructing
cuckoo hashing tables. We will describe each algorithm and then analyze the
construction times.

We quickly recall the most common usage of cuckoo hashing in cryptographic
primitives. Typically, one party will have an input set of items $X$ and will
construct a cuckoo hashing table based on the set $X$ (or an encrypted version
of the set $X$). In particular, the construction of the cuckoo hashing table
is done locally by the party and the transcript is typically not revealed to
any other parties including the adversary. Therefore, in this section, we will
consider the running time of the construction algorithm as if it were being
performed by a single party without any privacy considerations. In that case
that one is concerned about timing side channels, the construction algorithm
may be padded to execute in the worst case running time.

Given the above, our goal is to bound the time of construction algorithms
where all $n$ items are given as input. In other words, we are going to bound
the time of inserting all $n$ items. In contrast, most previous works
considered the time to bound the insertion of a single item.

For convenience, we consider the most important setting for cuckoo hashing
with a large number of hash functions $k$ where stashes are empty $s=0$,
entries may store at most one item $\ell=1$ and the number of entries is
linear in the number of items $b=O(n)$. We will consider an arbitrary number
of hash functions where $k$ is a sufficiently large constant.111We restrict
$k$ to be a sufficiently large constant to ensure that both random walks and
BFS behave well. For small $k$, these algorithms may behave differently. For
example, random walks with $k=2$ are essentially following a fixed path. The
analysis can be extended to consider arbitrary stash sizes $s$, entry sizes
$\ell$ and number of entries $b$.

Finally, we note that a lot of our analysis will rely on prior works that
analyze these insertion/construction algorithms for cuckoo hashing. Most prior
works consider $k=O(1)$. For example, the bounds for breadth first search in
[FPSS05] proved $O(1)$ insertion time. However, the real bound is
$O(\mathsf{poly}(k))$ that is no longer constant for super-constant $k$. In
our work, we modify the proofs of prior works to obtain bounds with respect to
arbitrary $k$.

Construction Algorithm Expected Time Worst Case Time with $1-\mathsf{negl}(n)$
Probability Maximum Cardinality Bipartite Matching [CKL+22]
$\tilde{O}((nk)^{1+o(1)})$ $\tilde{O}((nk)^{1+o(1)})$ Breadth First Search
[FPSS05] $O(n\cdot\mathsf{poly}(k))$ $O(n^{1+\alpha}\cdot k)$ Breadth First
Search, $k=\omega(\log n)$ $O(nk)$ $O(nk)$ Random Walk [FMM09, FPS13, FJ19,
Wal22] $O(nk)$ $O(nk\cdot\mathsf{polylog}(n))$ Local Search Allocation [Kho13]
$O(nk)$ $O(nk\cdot\mathsf{polylog}(n))$

Table 1: A comparison table of construction algorithms and their expected and
worst case running times. We consider the most important setting of
sufficiently large $k$, $s=0$, $\ell=1$ and $b=O(n)$.

###  F.1 Maximum Cardinality Bipartite Matching

Any allocation of the items corresponds to a matching in the bipartite graph
corresponding to the cuckoo hashing scheme. Therefore, one can re-phrase the
goal of cuckoo hashing allocation as determining whether the maximum
cardinality bipartite matching is as large as the number of left vertices.
Therefore, we can use any maximum cardinality bipartite matching algorithm
such as Ford and Fulkerson [FF56].

The most efficient one to date is due to Chen et al. [CKL+22] that runs in
$\tilde{O}(m^{1+o(1)})$ time for graphs with $m$ edges. In our setting, we
have that $m=nk$ to get that the worst case running time is
$\tilde{O}((nk)^{1+o(1)})$. However, we will analyze algorithms with better
guarantees below.

###  F.2 Breadth First Search

Description. Next, we consider breadth first search (BFS) that incrementally
inserts a new item to a current left perfect matching. A very convenient way
to maintain perfect left matchings corresponding to allocations is by using
the direction of edges. All edges outside of the current matching are directed
from left vertices to right vertices (prior works have used this technique
including [FPSS05]). On the other hand, all edges in the matching are directed
from right vertices to left vertices. When adding a new item (i.e., a left
vertex), the goal of inserting the item is equivalent to finding an
alternating path from the new vertex to any right vertex that is currently
unoccupied. By the orientation of the edges, this amounts to simply finding a
path from the new left vertex to any free right vertex. To obtain a
construction algorithm, we can run breadth first search for all $n$ items to
be allocated. Note, this is a perfect construction algorithm as a new item may
be inserted if and only if there exists a path to a free right vertex and BFS
will always find such a path if it exists.

Analysis. Recall that breadth first search (BFS) may be used to incrementally
insert new items into a cuckoo hash table. At a high level, BFS starts from
the node corresponding to the item in the cuckoo graph and attempts to find an
augmenting path to an empty entry node.

We will build upon the proof techniques of Fotakis, Pagh, Sanders and Spirakis
[FPSS05].222The paper considers cuckoo hashing with a single shared table.
However, their proofs rely on vertex expansion in the cuckoo graph that is
better when considering multiple disjoint tables. At a high level, they proved
that BFS takes $O(\mathsf{poly}(k))$ expected time and $O(n^{\alpha}\cdot k)$
worst case time for some constant $0<\alpha<1$ except with probability
$(2^{-\Omega(n)}+n^{-\Omega(k)})$.333The original theorem was meant to apply
for all values of $k\geq 8$. As a result, they could only obtain worst case
times except with $1-n^{-\Omega(1)}$ probability. We formally present the
theorem below.

######  Theorem 15 ([FPSS05]).

Let $k=\omega(1)$, $s=0$, $\ell=1$ and $b=O(n)$. For cuckoo hashing
$\mathsf{CH}(k,b,\ell,s)$ with a BFS construction algorithm where
$\mathcal{H}$ is $(nk)$-wise independent, then a single item will be inserted
using a random walk in expected time $O(\mathsf{poly}(k))$ and worst case time
$O(n^{\alpha}\cdot k)$ for some constant $0<\alpha<1$ except with probability
$1/2^{\Omega(n)}+1/n^{\Omega(k)}$.

######  Theorem 16.

Let $k=\omega(1)$, $s=0$, $\ell=1$ and $b=O(n)$. For cuckoo hashing
$\mathsf{CH}(k,b,\ell,s)$ with a BFS construction algorithm where
$\mathcal{H}$ is $(nk)$-wise independent, the following holds:

  * •

The expected construction time is $O(n\cdot\mathsf{poly}(k))$.

  * •

The worst case construction time is $O(n^{1+\alpha}\cdot k)$ except with
negligible probability.

###### Proof.

The expected time follows trivially from Theorem 15. As $k=\omega(1)$, the
worst case time of any insertion is $O(n^{\alpha}\cdot k)$ from Theorem 15
meaning the worst case time is $O(n^{1+\alpha}\cdot k)$. ∎

To our knowledge, it remains an open problem to prove worst case bounds for
BFS with high probability that are sub-polynomial.

We note that we can prove better bounds for BFS for larger values of
$k=\omega(\log n)$. In particular, we show that BFS will terminate after the
first step except with negligible probability. As a result, we can prove
optimal worst case construction times of $O(nk)$.

######  Theorem 17.

Let $k=\omega(\log n)$, $s=0$, $\ell=1$ and $b=O(n)$. For cuckoo hashing
$\mathsf{CH}(k,b,\ell,s)$ with a BFS construction algorithm where
$\mathcal{H}$ is $(nk)$-wise independent, the following holds:

  * •

The expected construction time is $O(nk)$.

  * •

The worst case construction time is $O(nk)$ except with negligible
probability.

###### Proof.

Consider the BFS algorithm at the first step. Assuming that $b=\alpha n$ for
some constant $\alpha>1$, we know that the hash table is only half full.
Suppose the number of entries in the first table is $n_{1}$, the number of
entries in the second table is $n_{2}$ and so forth such that
$n_{1}+\ldots+n_{k}\leq n$. The probability that the $k$ different entries are
all occupied is at most $(n_{1}/(\alpha n/k))\cdots(n_{k}/(\alpha n/k))$ that
is maximized by setting $n_{1}=\ldots=n_{k}=k/n$ to get an upper bound on the
probability of $(1/\alpha)^{k}$. As $k=\omega(\log n)$ and $\alpha$ is a
constant strictly greater than $1$, the probability that the first step of BFS
fails to insert the item is negligible. As a result, we get that the worst
case time to insert a single item is $O(k)$ except with negligible
probability. This implies the expected and worst case time (except with
negligible probability) to insert $n$ items is $O(nk)$. ∎

###  F.3 Random Walk

Description. The last insertion algorithm considers random walks. Using the
same idea of find alternating paths from BFS, we could instead try to find
such a path using a random walk. In general, random walks are considered more
efficient as they require less memory usage compared to BFS as one does not
need to remember the prior paths. To do this, as the random walk visits
entries, it will evict items from full entries and only remember that item
that it will try to place next.

The main challenge is that a random walk can, in theory, run forever without
terminating even if an alternating path exists. In the past, cuckoo hashing
with random walks will bound the length of the random walk before simply
terminating and putting the item in the stash. However, this does not work for
us if we wish to obtain a perfect construction algorithm.

Instead, we can obtain a perfect construction algorithm using a different
approach. In particular, we can let the random walk algorithm execute until
trying a path of length $O(n)$. If such a path has not yet been found, we can
instead execute a BFS to find the alternating path that also has worst case
$O(nk)$ overhead.

Worst Case Analysis. We use the results of Fountoulakis, Panagiotou and Steger
[FPS13] that consider the running time of inserting a single item using a
random walk. We could also use the results of Frieze, Melsted and Mitzenmacher
[FMM09] as we ignore additional log factors and consider larger values of $k$.
At a high level, they show a random walk must traverse $\mathsf{polylog}(n)$
nodes before finding an augmenting path.

######  Theorem 18 ([FPS13]).

Let $k\geq 3$, $s=0$, $\ell=1$ and $b=O(n)$. For cuckoo hashing
$\mathsf{CH}(k,b,\ell,s)$ where $\mathcal{H}$ is $(nk)$-wise independent, then
a single item will be inserted using a random walk in worst case time
$O(\mathsf{polylog}(n))$ except with probability $1/n$.444In the original
paper [FPS13], the authors showed this is true except with probability
$1/n^{\alpha}$ for $\alpha>0$. However, by increasing the running time by
poly-logarithmic factors, one can drive the probability down to $1/n$.

Using this result, we can immediately derive the following theorem that bounds
the expected and worst case running times of random walks.

######  Theorem 19.

Let $k\geq 3$, $s=0$, $\ell=1$ and $b=O(n)$. For cuckoo hashing
$\mathsf{CH}(k,b,\ell,s)$ with a random walk construction algorithm where
$\mathcal{H}$ is $(nk)$-wise independent, then the worst case construction
time is $O(nk\cdot\mathsf{polylog}(n))$ except with negligible probability.

###### Proof.

We note that we can consider random variables $X_{i}=1$ if and only if the
$i$-th item that is inserted exceeds time $O(\mathsf{polylog}(n))$. We know
that $\Pr[X_{i}=1]\leq 1-1/n$. If we let $X=X_{1}+\ldots+X_{n}$, we know that
$\Pr[X>\log^{2}n]\leq\mathsf{negl}(n)$ by Chernoff Bounds. Therefore, the
total construction time except with negligible probability is $(n-X)\cdot
O(\mathsf{polylog}(n))+X\cdot O(nk)=O(nk\cdot\mathsf{polylog}(n))$ as
$X\leq\log^{2}n$ except with negligible probability. ∎

Expected Analysis. We note the same papers [FMM09, FPS13] that we relied upon
above proved a $\mathsf{polylog}(n)$ expected time bound for random walks.
Instead, we can move to more recent works [FJ19, Wal22] that proved constant
bounds for the expected times of random walks. Again, they proved bounds for
constant $k$, but one can re-interpret their results for arbitrary $k$.

######  Theorem 20 ([Wal22]).

Let $k\geq 3$, $s=0$, $\ell=1$ and $b=O(n)$. For cuckoo hashing
$\mathsf{CH}(k,b,\ell,s)$ with a random walk construction algorithm where
$\mathcal{H}$ is $(nk)$-wise independent, then a single item will be inserted
using a random walk in expected time $O(k)$.

By linearity of expectation, the expected construction time for random walks
immediately follows as $O(nk)$.

###  F.4 Local Search Allocation

Description. Khosla [Kho13] presented a construction algorithm for the case of
$k\geq 3$, $\ell=1$, $s=0$ and $b=O(n)$. Unlike BFS and random walks, this is
a construction algorithm that takes the entire input at once and aims to
create a hash table.

At a high level, each of the $b$ entries maintains an integer label.
Initially, all labels are $0$. Whenever an item is inserted, it will search
amongst its $k$ options to find the minimum label entry. The item will be
placed into that entry. Furthermore, the label of that entry is increased to
be one more than the minimum amongst the other $k-1$ options for that item. If
that entry was already occupied, we repeat this algorithm by evicting the
original item in the entry. This is considered a construction algorithm as one
must maintain the labels throughout the construction. If one is content with
storing all labels, it is possible to convert this algorithm into an insertion
algorithm as well.

In comparison with BFS, we note that both algorithms require linear storage.
This algorithm provides stronger expected and worst-case guarantees compared
to BFS. Compared to random walk, this algorithm requires more storage but can
still provide stronger worst case guarantees. Furthermore, experiments in the
original paper [Kho13] showed that the algorithm outperforms random walks in
terms of running time.

Analysis. In the paper [Kho13], it was proven that the construction of this
algorithm is $O(nk)$ except with $1/\mathsf{poly}(n)$ probability. As a
result, we can immediately get that the expected running time is also $O(nk)$.
Similarly, it was shown that the maximum contribution to the total running
time by any vertex (i.e., the label) is at most $O(\log n)$ with
$1/\mathsf{poly}(n)$ probability. As a result, we can obtain a worst case
bound as well by showing that at most $\mathsf{polylog}(n)$ such vertices will
exceed $O(\log n)$ except with negligible probability.

###  F.5 Robust Construction Algorithms

In the prior sections, we only analyze the running times for construction
algorithms with the assumption that items to be constructed are chosen
independent of the hash functions. This is necessary as a computationally
unbounded adversary can find sequences that incur the worst case construction
times for random walks and BFS. In fact, such powerful adversaries can find
sets of items that will cause the construction algorithm to fail. If one
wishes to consider cuckoo hashing construction with respect to compromised
randomness known by an adversary, then the best option is to use deterministic
algorithms. To our knowledge, the best deterministic algorithm for maximum
bipartite matching is $(nk)^{4/3+o(1)}$ by Kathuria, Liu and Sidford [KLS22].