D ifferentiate E verything with a R eversible E mbeded D omain -S pecific L anguage
Anonymous authors
Paper under double-blind review
A bstract
Reverse-mode automatic di ﬀ erentiation (AD) su ﬀ ers from the issue of having too much space overhead to trace back intermediate computational states for back-propagation.
The traditional method to trace back states is called checkpointing that stores intermediate states into a global stack and restore state through either stack pop or re-computing.
The overhead of stack manipulations and re-computing makes the general purposed (not tensor-based) AD engines unable to meet many industrial needs.
Instead of checkpointing, we propose to use reverse computing to trace back states by designing and implementing a reversible programming eDSL, where a program can be executed bi-directionally without implicit stack operations.
The absence of implicit stack operations makes the program compatible with existing compiler features, including utilizing existing optimization passes and compiling the code as GPU kernels.
We implement AD for sparse matrix operations and some machine learning applications to show that our framework has the state-of-the-art performance.
I ntroduction
Most of the popular automatic di ﬀ erentiation (AD) tools in the market, such as TensorFlow (Abadi et al., 2015), Pytorch (Paszke et al., 2017), and Flux (Innes et al., 2018) implements reverse mode
AD at the tensor level to meet the need in machine learning.
Later, People in the scientiﬁc computing domain also realized the power of these AD tools, they use these tools to solve scientiﬁc problems such as seismic inversion (Zhu et al., 2020), variational quantum circuits simulation (Bergholm et al., 2018;
Luo et al., 2019) and variational tensor network simulation (Liao et al., 2019; Roberts et al., 2019).
To meet the diverse need in these applications, one sometimes has to deﬁne backward rules manually, for example
To di ﬀ erentiate sparse matrix operations used in Hamiltonian engineering (Hao Xie & Wang)
, people deﬁned backward rules for sparse matrix multiplication and dominant eigensolvers (Golub & Van Loan, 2012), 2.
In tensor network algorithms to study the phase transition problem (Liao et al., 2019; Seeger et al., 2017; Wan & Zhang, 2019; Hubig, 2019), people deﬁned backward rules for singular value decomposition (SVD) function and QR decomposition (Golub & Van Loan, 2012).
Instead of deﬁning backward rules manually, one can also use a general purposed AD (GP-AD) framework like Tapenade (Hascoet & Pascual, 2013), OpenAD (Utke et al., 2008) and
Zygote (Innes, 2018; Innes et al., 2019).
Researchers have used these tools in practical applications such as bundle adjustment (Shen & Dai, 2018) and earth system simulation
(Forget et al., 2015), where di ﬀ erentiating scalar operations is important.
However, the power of these tools are often limited by their relatively poor performance.
In many practical applications, a program might do billions of computations.
In each computational step, the AD engine might cache some data for backpropagation.
(Griewank & Walther, 2008)
Frequent caching of data slows down the program signiﬁcantly, while the memory usage will become a bottleneck as well.
Caching implicitly also make these frameworks incompatible with kernel functions.
To avoid such issues, we need a new
GP-AD framework that does not cache automatically for users.
In this paper, we propose to implement the reverse mode AD on a reversible (domain-speciﬁc) programming language (Perumalla, 2013; Frank, 2017), where intermediate states can be traced backward without accessing an implicit stack.
Reversible programming allows people to utilize the reversibility to reverse a program.
In machine learning, reversibility is proven to substantially decrease the memory usage in unitary recurrent neural networks (MacKay et al., 2018), normalizing ﬂow (Dinh et al., 2014), hyper-parameter learning (Maclaurin et al., 2015) and residual neural networks (Gomez et al., 2017; Behrmann et al., 2018).
Reversible programming will make these happen naturally.
The power of reversible programming is not limited to handling these reversible applications, any program can be written in a reversible style.
Converting an irreversible program to the reversible form would cost overheads in time and space.
Reversible programming provides a ﬂexible time-space trade-o ﬀ scheme that di ﬀ erent with checkpointing (Griewank, 1992; Griewank & Walther, 2008; Chen et al., 2016), reverse computing (Bennett, 1989; Levine & Sherman, 1990), to let user handle these overheads explicitly.
There have been many prototypes of reversible languages like Janus (Lutz, 1986), R (not the popular one) (Frank, 1997), Erlang (Lanese et al., 2018) and object-oriented ROOPL
(Haulund, 2017).
In the past, the primary motivation to study reversible programming is to support reversible computing devices (Frank &
Knight
Jr, 1999) such as adiabatic complementary metal-oxide-semiconductor (CMOS) (Koller & Athas, 1992), molecular mechanical computing system (Merkle et al., 2018) and superconducting system (Likharev, 1977; Semenov et al., 2003;
Takeuchi et al., 2014; 2017), and these reversible computing devices are orders more energy-e ﬃ cient.
Landauer proves that only when a device does not erase information (i.e. reversible), its energy e ﬃ ciency can go beyond the thermal dynamic limit.
(Landauer, 1961; Reeb & Wolf, 2014)
However, these reversible programming languages can not be used directly in real scientiﬁc computing, since most of them do not have basic elements like ﬂoating point numbers, arrays, and complex numbers.
This motivates us to build a new embedded domain-speciﬁc language (eDSL) in Julia (Bezanson et al., 2012; 2017) as a new playground of GP-AD.
In this paper, we ﬁrst introduce the language design of NiLang in Sec. 3.
In Sec. 4, we explain the implementation of automatic di ﬀ erentiation in NiLang.
In Sec. 5, we benchmark the performance of NiLang’s AD with other AD software and explain why it is fast.
R everse computing as an A lternative of C heckpointing
Reverse computing and checkpointing share many similarities.
Consider an irreversible program with pure forward computing time T and run-time memory S .
Table 1 shows that both schemes su ﬀ er from a space overhead that linear to time when executing in time
O ( T ).
In reversible computing, to reach a minimum space overhead of O ( S log( T / S ))
(Bennett, 1989; Levine &
Sherman, 1990; Perumalla, 2013), the program has a polynomial overhead in time.
While in checkpointing, there can be no space overhead.
Since one can just recompute from beginning to obtain any intermediate state with time complexity O ( T 2 ).
[Table]
In practical using cases, we need to trade o ﬀ space and time.
The most successful checkpointing algorithm that is widely used in automatic di ﬀ erentiation is the treeverse algorithm in Fig. 1(a).
Where the computational process is binomially partitioned into d sectors.
At the beginning of each sector, a snapshot is stored in the main memory.
The states in the last sector are retrieved by the above space-e ﬃ cient O
( T 2 ) algorithm.
After that, the last checkpoint can be freed.
The remaining sectors are further partitioned into d − l + 1 sub-sectors, where l is the sector index counting from the tail.
The earlier sectors have more quota of snapshots while the latter sectors have less so that the total number of snapshots remain the same.
Recursively apply this treeverse algorithm t times
[Figure]
P
Q Q
P P
P
[Figure]
until the sector size is 1.
The approximated overhead in time and space are
[Equation]
Where T = η
( t , d ).
By carefully choosing either a t or d , the overhead in time and space can be both logarithmic.
On the other side, Bennett’s trade-o ﬀ of reverse computing also has a recursive structure as shown in Fig.
1 (b).
But the program is evenly partitioned into k sectors.
The program marches forward ( P process) for k steps to obtain the ﬁnal state s k + 1 , then backward ( Q process) from the k − 1th step to erase the states in between s 1 < i ≤ k .
Each sector is further divided in to k sub-sectors and recursive run the above compute-copy-uncompute process.
The time and space complexities are
[Equation]
Here, the overhead in time is polynomial, which is slightly worse than the treeverse algorithm.
Treeverse like partition does not apply here because one can not complete the ﬁrst sweep to create initial checkpoints without introducing any space overheads in reversible computing.
The pseudocode of Bennett’s time-space trade-o ﬀ algorithm is shown in Listing.
The ﬁrst argument { s 1 ,...} is the collection of states, k is the number of partitions, i and len are the starting point and length of the working sector.
A function call changes variables inplace.
“ ∼ ” is the symbol of uncomputing, which means undoing a function call.
Statement s i + 1 ← 0 allocates a zero state and add it to the state collection.
Its inverse s i + 1 → 0 discards a zero cleared state from the collection.
The NiLang implementation is in Appendix A.
Another interesting observation is reversible computing is not always more energy e ﬃ cient irreversible computing. The time to uncompute a unit of memory is exponential to n as Q n = (2 k − 1) n , and the computing energy also increases exponentially.
On the other side, the amount of energy to erase a memory unit is constant.
When (2 k − 1) n > 1 /ξ , erasing the memory irreversibly is more energy-e ﬃ cient, where ξ is the energy ratio between a reversible operation (an instruction or a gate) and its irreversible counterpart. For this reason, we think a reversible programming language is more suitable as an eDSL inside an irreversible host language rather than a stand-alone one. One can use the reversible operations at the microscopic level like registers and L1 cache level to save energy, and use the existing irreversible software stacks at the macroscopic level like the main memory level to discard information directly.
[Table]
The reverse computing does not show any advantage in worse-case complexity comparing with checkpointing.
But the following traits make it perform better in many practical applications.
First , it can make use of the reversibility to save memory.
The above discussion assumes every operation is irreversible, however, the most program contains a lot of reversible instructions like the accumulation multiplication operations in many BLAS functions and sparse matrix functions.
In
Appendix B.2, we show how to implement a unitary matrix multiplication without introducing overheads in space and time.
Second , reverse computing does not allocate automatically for users, user can optimize the memory access patterns for their own devices like GPU.
Third , reverse computing can operate on e ﬀ ective codes, so that it ﬁts better with modern languages.
In
Appendix B.1, we show how to manipulate inplace functions on arrays with NiLang.
Fourth , reverse computing can utilize the existing compiler to optimize the code better because it does not automatically introduce global stack operations that harm the purity of functions.
Fifth , reverse computing encourages the user to think reversibly.
Reversible thinking can lead the user to a constant memory, constant time overhead implementation of chained multiplication algorithms as shown in Appendix B.3.
Transpiling a regular code to a reversible code is not hard, but it is unlikely to provide the user with better performance than optimal checkpointing.
Instead, thinking about how to write instructions and control ﬂows reversibly does.
L anguage design
NiLang is an embedded domain-speciﬁc language (eDSL)
NiLang built on top of the host language
Julia (Bezanson et al., 2012; 2017).
Julia is a popular language for scientiﬁc programming and machine learning.
We choose Julia mainly for speed.
Julia is a language with high abstraction, however, its clever design of type inference and just in time compiling make it has a C like speed.
Meanwhile, it has rich features for meta-programming.
Its package for pattern matching MLStyle allows us to deﬁne an eDSL in less than 2000 lines.
Comparing with a regular reversible programming language, NiLang features array operations, rich number systems including ﬂoating-point numbers, complex numbers, ﬁxed-point numbers, and logarithmic numbers.
It also implements the compute-copy-uncompute (Bennett, 1973) macro to increase code reusability.
Besides the above reversible hardware compatible features, it also has some reversible hardware incompatible features to meet the practical needs.
For example, it views the ﬂoating-point + and − operations as reversible.
It also allows users to extend instruction sets and sometimes inserting external statements.
These features are not compatible with future reversible hardware.
NiLang’s source code is available online, we will put a link here after the anonymous open review session.
By the time of writing, the version of NiLang is v0.7.3.
3.1 R eversible functions and instructions
Mathematically, any irreversible mapping y = f(args...) can be trivially transformed to its reversible form y += f(args...) or y (cid:89) = f(args...)
( (cid:89) is the bit-wise XOR ), where y is a pre-emptied variable.
But in numeric computing with ﬁnite precision, this is not always true.
The reversibility of arithmetic instruction is closely related to the number system.
For integer and ﬁxed point number system, y += f(args...) and y -= f(args...) are rigorously reversible.
For logarithmic number system and tropical number system (Speyer & Sturmfels, 2009), y *= f(args...) and y /= f(args...)
as reversible (not introducing the zero element).
While for ﬂoating point numbers, none of the above operations are rigorously reversible.
However, for convenience, we ignore the round-o ﬀ errors in ﬂoating-point + and - operations and treat them on equal footing with ﬁxed-point numbers in the following discussion. In Appendix ??
, we will show doing this is safe in most cases provided careful implementation.
Other reversible operations includes SWAP , ROT , NEG et.
al., and this instruction set is extensible.
One can deﬁne a reversible multiplier in NiLang as in Listing. 2.
Listing 2: A reversible multiplier julia> [Table] NiLang [Table]
end [Table] [Table] [Table]
Macro @i generates two functions that are reversible to each other, multiplier and ∼ multiplier , each deﬁnes a mapping R 3 → R 3 .
The !
after a symbol is a part of the name, as a conversion to indicate the mutated variables.
3.2 R eversible memory management
A distinct feature of reversible memory management is that the content of a variable must be known when it is deallocated.
We denote the allocation of a pre-emptied memory as x ← 0 , and its inverse, deallocating a zero emptied variable, as x → 0 .
An unknown variable can not be deallocate, but can be pushed to a stack pop out later in the uncomputing stage.
If a variable is allocated and deallocated in the local scope, we call it an ancilla.
Listing.
3 deﬁnes the complex valued accumulative log function.
Listing 3: Reversible complex valued log function y += log( | x | )
+ i
Arg( x ).
[Figure]
end
Listing 4: Compute-copy-uncompute version of Listing.@i
[Table] [Table]
~
[Table]
Here, the macro @inline tells the compiler that this function can be inlined.
One can input “ ← ” and “ → ” by typing “ \ leftarrow[TAB KEY]” and “ \ rightarrow[TAB KEY]” respectively in a Julia editor or REPL.
NiLang does not have immutable structs, so that the real part y!.re and imaginary y!.im of a complex number can be changed directly.
It is easy to verify that the bottom two lines in the function body are the inverse of the top two lines.
i.e., the bottom two lines uncomputes the top two lines.
The motivation of uncomputing is to zero clear the contents in ancilla n so that it can be deallocated correctly.
Compute-copy-uncompute is a useful design pattern in reversible programming so that we created a pair of macros @routine and ∼ @routine for it.
One can rewrite the above function as in Listing. 4.
3.3 R eversible control flows
One can deﬁne reversible if , for and while statements in a reversible program.
Fig.
2 (a) shows the ﬂow chart of executing the reversible if statement.
There are two condition expressions in [Figure]
stmts
[Figure]
T F error post post stmts
F
T T
F error pre
T error
F pre equals to post
[Figure]
this chart, a precondition and a postcondition.
The precondition decides which branch to enter in the forward execution, while the postcondition decides which branch to enter in the backward execution.
The pseudo-code for the forward and backward passes are shown in Listing.
5 and Listing.
Listing
: Translating a reversible if statement (forward)
[Figure]
Listing
: Translating a reversible if statement (backward)
[Figure]
Fig.
2 (b) shows the ﬂow chart of the reversible while statement.
It also has two condition expressions.
Before executing the condition expressions, the program presumes the postcondition is false.
After each iteration, the program asserts the postcondition to be true.
To reverse this statement, one can exchange the precondition and postcondition, and reverse the body statements.
The pseudo-code for the forward and backward passes are shown in Listing.
7 and Listing.
Listing 7: Translating a reversible while statement (forward)
[Figure]
end
Listing
: Translating a reversible while statement (backward)
[Figure]
end
The reversible for statement is similar to the irreversible one except that after execution, the program will assert the iterator to be unchanged.
To reverse this statement, one can exchange start and stop and inverse the sign of step .
Listing.
9 computes the Fibonacci number recursively and reversibly.
Listing 9: Computing Fibonacci number recursively and reversibly.
@i function rrfib(out!, n) @invcheckoff if (n >= 1, ~) counter ← 0 counter += n while (counter > 1, counter!=n) rrfib(out!, counter-1) counter -= 2 end counter
[Table]
= n %
[Table]
→end
[Table]
= 1 end Here, out! is an integer initialized to 0 for storing outputs.
The precondition and postcondition are wrapped into a tuple.
In the if statement, the postcondition is the same as the precondition, hence we omit the postcondition by inserting a " ∼ " in the second ﬁeld for “copying the precondition in this ﬁeld as the postcondition”.
In the while statement, the postcondition is true only for the initial loop.
Once code is proven correct, one can turn o ﬀ the reversibility check by adding @invcheckoff before a statement.
This will remove the reversibility check and make the code faster and compatible with GPU kernels (kernel functions can not handle exceptions).
R eversible automatic differentiation 4.1 B ack propagation
[Figure]
We decompose the problem of reverse mode AD into two sub-problems, reversing the code and computing ∂ [single input] ∂ [multiple outputs] .
Reversing the code is trivial in reversible programming, while the second sub-problem is similar to forward mode automatic di ﬀ erentiation that computes ∂ [multiple outputs] ∂
[single input] .
Inspired by the Julia package ForwardDi ﬀ (Revels et al., 2016), we use operator overloading rather than the source to source transformation for better extensibility.
We wrap each output variable with a composite type GVar that containing an extra gradient ﬁeld, and feed it into the reversed generic program.
Instructions are multiple dispatched to corresponding gradient instructions that update the gradient ﬁeld of GVar at the same time.
By reversing this gradient program, we can obtain the gradient program for the reversed program R too.
One can deﬁne the adjoint (“adjoint” here means a reversed program updating gradients) of a primitive instruction as a reversible function on either the function itself or its reverse since the adjoint of a function’s reverse is equivalent to the reverse of the function’s adjoint.
[Equation]
It can be easily veriﬁed by applying the above two mappings consecutively, which turns out to be an identity mapping considering
[Equation]
The implementation details are described in Appendix C. In most languages, operator overloading brings signiﬁcant overheads of function calls and object allocation and deallocation.
But in a language with type inference and just in time compiling like Julia, the boundary between two approaches are vague.
The compiler inlines small functions, packs an array of constant sized immutable objects into a continuous memory, and truncates unnecessary branches automatically.
4.2 H essians
Combining forward mode AD and reverse mode AD is a simple yet e ﬃ cient way to obtain Hessians.
By wrapping the elementary type with Dual deﬁned in package ForwardDi ﬀ and throwing it into the gradient program deﬁned in NiLang, one obtains one row / column of the Hessian matrix.
We will use this approach to compute Hessians in the graph embedding benchmark in Sec. D.2.
4.3 CUDA kernels
CUDA programming is playing a signiﬁcant role in high-performance computing.
In Julia, one can write GPU compatible functions in native Julia language with KernelAbstractions.
(Besard et al., 2017)
Since NiLang does not push variables into stack automatically for users, it is safe to write di ﬀ erentiable GPU kernels with NiLang.
We will di ﬀ erentiate CUDA kernels with no more than extra 10 lines in the bundle adjustment benchmark in Sec. 5.
B enchmarks
We benchmark our framework with the state-of-the-art GP-AD frameworks, including source code transformation based Tapenade and Zygote and operator overloading based ForwardDi ﬀ and
ReverseDi ﬀ .
Since most tensor based AD software like famous TensorFlow and PyTorch are not designed for the using cases used in our benchmarks, we do not include those package to avoid an unfair comparison.
In the following benchmarks, the CPU device is Intel(R)
Xeon(R) Gold
CPU @ 2.10GHz, and the GPU device is NVIDIA Titan V. For NiLang benchmarks, we have turned the reversibility check o ﬀ to achieve a better performance.
We reproduced the benchmarks for Gaussian mixture model (GMM) and bundle adjustment in
Srajer et al.
(2018) by re-writing the programs in a reversible style.
We show the results in Fig.
The Tapenade data is obtained by executing the docker ﬁle provided by the original benchmark, which provides a baseline for comparison.
[Figure]
NiLang’s objective function is ∼ 2 × slower than normal code due to the uncomputing overhead.
In this case, NiLang does not show advantage to Tapenade in obtaining gradients, the ratio between computing the gradients and the objective function are close.
This is because the bottleneck of this model is the matrix vector multiplication, traditional AD can already handle this function well.
The extra memory used to reverse the program is negligible comparing to the original program as shown [Figure] in Fig. 4.
The backward pass is not shown here, it is just two times the reversible program in order to store gradients.
The data is obtained by counting the main memory allocations in the program manually.
The analytical expression of memory usage in unit of ﬂoating point number is
[Equation]
where d and k are the size and number of covariance matrices.
P =
d ( d + 1) 2 k + k + dk is the size of parameter space.
The memory of the dataset ( d × N ) is not included because it will scale as N .
Due to the hardness of estimating peak memory usage, the Tapenade data is missing here.
The ForwardDi ﬀ memory usage is approximately the original size times the batch size, where the batch size is 12 by default.
In the bundle adjustment benchmark, NiLang performs the best on CPU.
We also compiled our adjoint program to GPU with no more than 10 lines of code with KernelAbstractions, which provides another ∼ 200 × speed up.
Parallelizing the adjoint code requires the forward code not reading the same variable simultaneously in di ﬀ erent threads, and this requirement is satisﬁed here.
The peak memory of the original program and the reversible program are both equal to the size of parameter space because all “allocation”s happen on registers in this application.
One can ﬁnd more benchmarks in Appendix D, including di ﬀ erentiating sparse matrix dot product and obtaining Hessians in the graph embedding application.
D iscussion
In this work, we demonstrate a new approach to back propagates a program called reversible programming AD by designing a reversible eDSL
NiLang.
NiLang is a powerful tool to di ﬀ erentiate code from the source code level so that can be directly useful to machine learning researches.
It can generate e ﬃ cient backward rules, which is exempliﬁed in Appendix E.
It can also be used to di ﬀ erentiate reversible neural networks like normalizing ﬂows (Kobyzev et al., 2019) to save memory, e.g. back-propagating NICE network (Dinh et al., 2014) with only constant space overheads.
Except for the above immediate beneﬁts to the machine learning ecosystem, we also want to emphasize the broader impact of AD on reverse computing.
AD on reverse computing is related to two of the central issues limiting many machine learning applications.
One is how to overcome the memory wall that has been extensively discussed in the main text, and another is how to reduce energy consumption.
AD on reverse computing provides a software stack that connecting machine learning with further technologies to save energy.
A cknowledgments
The authors are grateful to the people who help improve this work and fundings that sponsored the research.
To meet the anonymous criteria, we will add the acknowledgments after the open review session
A N i L ang implementation of B ennett ’ s time - space trade - off algorithm
using NiLang, Test
PROG_COUNTER = Ref(0) #
(2k-1)^n PEAK_MEM = Ref(0)
# n*(k-1)+ @i function bennett(f::AbstractVector, state:: [Table] {Int,T}, k::Int, [Table]
, len) where T if (len == 1, ~) state [Table] +1] ← zero(T) f[base](state
[Table]
], state[ [Table] ]) @safe [Table]
[Table] [Table] [Table] [Table]
[Table]
[] =
[Table]
state))
[Table]
else n ←
[Table] = [Table] # the P process for j=1:
k bennett(f, state, k,
[Table]
+n*(j
[Table]
1), n) end # the Q process for j= [Table] [Table] ~bennett(f, state, k,
[Table] +n*(j-1), n) end n [Table] =
[Table]
n
[Table]
end end k
= 4 n = N =
[Table]
n state = [Table] (1=> [Table]
f [Table] x) = [Table]
[Table] [Table] instructions = fill(PlusEq(f), N) # run the program [Table]
[Table] [Table] [Table] , 1, N) @
[Table] [Table] [Table] [Table]
[Table] + [Table] [Table]
[] = [Table] [Table] [Table]
[Table]
n The input f is a vector of functions and state is a dictionary.
We also added some irreversible external statements (those marked with @safe ) to help analyse to program.
B C ases where reverse computing shows advantage B.1 H andling effective codes
Reverse computing can handling e ﬀ ective codes with mutable structures and arrays.
For example, the a ﬃ ne transformation can be implemented without any overhead.
Listing 11: Inplace a ﬃ ne transformation.
[Table]
end end
[Table]
end end
Here, the expression following the @safe macro is an external irreversible statement.
B.2 U tilizing reversibility
Reverse computing can utilize reversibility to trace back states without extra memory cost.
For example, we can deﬁne the unitary matrix multiplication that can be used in a type of memorye ﬃ cient recurrent neural network.
(Jing et al., 2016)
Listing 12: Two level decomposition of a unitary matrix.
[Figure]
B.3 E ncourages reversible thinking
Last but not least, reversible programming encourages users to code in a memory friendly style.
Since allocations in reversible programming are explicit, programmers have the ﬂexibility to control how to allocate memory and which number system to use.
For example, to compute the power of a positive ﬁxed-point number and an integer, one can easily write irreversible code as in Listing.
Listing 13: A regular power function.
[Equation]
Listing 14: A reversible power function.
[Equation]
Since the ﬁxed-point number is not reversible under *= , naive checkpointing would require stack operations inside a loop.
With reversible thinking, we can convert the ﬁxed-point number to logarithmic numbers to utilize the reversibility of *= as shown in Listing.
Here, the algorithm to convert a regular ﬁxed-point number to a logarithmic number can be e ﬃ cient.
(Turner, 2010)
C I mplementation of AD in N i L ang
To backpropagate the program, we ﬁrst reverse the code through source code transformation and then insert the gradient code through operator overloading.
If we inline all the functions in Listing.
4, the function body would be like Listing.
The automatically generated inverse program (i.e.
( y , x ) → ( y − log( x ) , x )) would be like Listing.
Listing 15: The inlined function body of
Listing.
(cid:7) [Equation]
[Figure]
[Equation]
Listing
: The inverse of Listing.
[Equation] [Figure]
[Equation]
To compute the adjoint of the computational process in Listing.
15, one simply insert the gradient code into its inverse in Listing.
The resulting inlined code is show in Listing.
Listing 17: Insert the gradient code into Listing.
16, the original computational processes are highlighted in yellow background.
[Figure]
[Equation]
[Figure]
[Table]
Here, @zeros TYPE var1 var2... is the macro to allocate multiple variables of the same type.
Its inverse operations starts with ∼ @zeros deallocates zero emptied variables.
In practice, “inserting gradients” is not achieved by source code transformation, but by operator overloading.
We change the element type to a composite type GVar with two ﬁelds, value x and gradient g .
With multiple dispatching primitive instructions on this new type, values and gradients can be updated simultaneously.
Although the code looks much longer, the computing time (with reversibility check closed) is not.
[Table]
The performance is unreasonably good because the generated Julia code is further compiled to
LLVM so that it can enjoy existing optimization passes.
For example, the optimization passes can ﬁnd out that for an irreversible device, uncomputing local variables n and nsq to zeros does not a ﬀ ect return values, so that it will ignore the uncomputing process automatically.
Unlike checkpointing based approaches that focus a lot in the optimization of data caching on a global stack, NiLang does not have any optimization pass in itself.
Instead, it throws itself to existing optimization passes in
Julia.
Without accessing the global stack, NiLang’s code is quite friendly to optimization passes.
In this case, we also see the boundary between source code transformation and operator overloading can be vague in a Julia, in that the generated code can be very di ﬀ erent from how it looks.
The joint functions for primitive instructions (:
+=)(sqrt) and (:-=)(sqrt) used above can be deﬁned as in Listing.
Listing
(cid:4)
@i @inline function (:
=)(
)(out
) where T @routine @invcheckoff begin @zeros T a b a += sqrt(x.x) b += 2 * a end
out!.x -= a x.g += out!.g / b
~@routine end (cid:6) (cid:5) D M ore B enchmarks D.1 S parse matrices
We benchmarked the call, uncall and backward propagation time used for sparse matrix dot product and matrix multiplication.
Here, we estimate the time for back propagating gradients rather than including both forward and backward
[Table]
since mul!
does not output a scalar as loss
[Table]
The time used for computing backward pass is approximately 1.5-3 times the Julia’s native forward pass.
This is because the instruction length of di ﬀ erentiating basic arithmetic instructions is longer than pure computing.
D.2 G raph embedding problem
Graph embedding can be used to ﬁnd a proper representation for an order parameter (Takahashi & Sandvik, 2020) in condensed matter physics.
People want to ﬁnd a minimum Euclidean space dimension k that a Petersen graph can embed into, that the distances between pairs of connected vertices are l 1 , and the distance between pairs of disconnected vertices are l 2 , where l 2 > l 1 .
The
Petersen graph is shown in Fig.
Let us denote the set of connected and disconnected vertex pairs as L 1 and L 2 , respectively.
This problem can be variationally solved with the following loss.
[Equation]
[Equation] The ﬁrst line is a summation of distance variances in two sets of vertex pairs, where Var( X ) is the variance of samples in X .
The second line is used to guarantee l 2 > l 1 , where X means taking the average of samples in X .
Its reversible implementation could be found in our benchmark repository.
We repeat the training for dimension k from 1 to 10.
In each training, we ﬁx two of the vertices and optimize the positions of the rest.
Otherwise, the program will ﬁnd the trivial solution with overlapped vertices.
For k < 5, the loss is always much higher than 0, while for k ≥ 5, we can get a loss close to machine precision with high probability.
From the k = 5 solution, it is easy to see l 2 / l 1 = √ 2.
An Adam optimizer with a learning rate 0 .
01 (Kingma & Ba) requires ∼ 2000 steps training.
The trust region Newton’s method converges much faster, which requires ∼ 20 computations of
Hessians to reach convergence.
Although training time is comparable, the converged precision of the later is much better.
Since one can combine ForwardDi ﬀ and NiLang to obtain Hessians, it is interesting to see how much performance we can get in di ﬀ erentiating the graph embedding program.
In Table 3, we show the the performance of di ﬀ erent implementations by varying the dimension k .
The number of parameters is 10 k .
As the baseline, (a) shows the time for computing the function call.
We have reversible and irreversible implementations, where the reversible program is slower than the irreversible native Julia program by a factor of ∼ 2 due to the uncomputing overhead.
The reversible program shows the advantage of obtaining gradients when the dimension k ≥ 3.
The larger the number of inputs, the more advantage it shows due to the overhead proportional to input size in forward mode AD.
The same reason applies to computing Hessians, where the combo of
NiLang and ForwardDi ﬀ gives the best performance for k ≥ 3. [Table]
embedding program.
k is the embedding dimension, the number of parameters is 10 k .
E P orting N i L ang to Z ygote
Zygote is a popular machine learning package in Julia.
We can port NiLang’s automatically generated backward rules to Zygote to accelerate some performance-critical functions.
The following example shows how to speed up the backward propagation of norm by ∼ 50 times.
Listing 20:
Porting NiLang to Zygote.
(cid:7) (cid:4)
[Table]
[Equation]
We ﬁrst import the norm function from Julia standard library LinearAlgebra.
Zygote’s builtin AD engine will generate a slow code and memory allocation of 339KB.
Then we write a reversible norm function r_norm with NiLang and port the backward function to Zygote by specifying the backward rule (the function marked with macro Zygote.@adjoint ).
Except for the speed up in computing time, the memory allocation also decreases to 23KB, which is equal to the sum of the original x and the array used in backpropagation.
[Equation]
The later one has a doubled size because GVar has an extra gradient ﬁeld. F E rror A nalysis
Running reversible programming with the ﬂoating pointing number system can introduce round-o ﬀ errors and make the program not reversible.
The quantify the e ﬀ ects, we use the leapfrog integrator to compute the orbitals of planets in our solar system as a benchmark.
Since the dynamics of our solar system are symplectic and the leapfrog integrator is time-reversible, the reversible program does not have overheads and the evolution time can go arbitrarily long with constant memory.
We compare the mean error in the ﬁnal axes of the planets and show the results in Fig.
Errors are computed by comparing with the results computed with high precision ﬂoating-point numbers.
One of the key steps that introduce round-o ﬀ error is the computation of acceleration .
If it is implemented as in Listing.
21, the round-o ﬀ error does not bring additional e ﬀ ect in the reversible context, hence we see overlapping lines “(regular)” and “(reversible)” in the ﬁgure.
This is because,
[Figure]
when returning a dirty (not exactly zero cleared due to the ﬂoating-point round-o ﬀ error) ancilla to the ancilla pool, the small remaining value will be zero-cleared automatically in NiLang.
The acceleration function can also be implemented as in Listing.
22, where the same variable rb is repeatedly used for compute and uncompute, the error will accumulate on this variable.
In both
Float64 (double precision ﬂoating point) and Float32 (single precision ﬂoating point) benchmarks, the results show a much lower precision.
Hence, simulating reversible programming with ﬂoating-point numbers does not necessarily make the results less reliable if one can avoid cumulative errors in the implementation.
Listing 21: Compute-uncompute on a new space rc (cid:7)
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
where
[Figure]
[Equation]
end (cid:6) (cid:5)
Listing 22: Compute-uncompute repeatedly on the same variable rb .
[Equation]
[Equation]
[Equation]
[Table]
[Equation]
end (cid:6)
[Equation]