Differentiate Everything with a Reversible Domain-Specific Language

1 Introduction Most of the popular automatic differentiation ( 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 scientific computing domain also realized the power of these AD tools , they use these tools to solve scientific 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 define backward rules manually , for example 1 . To differentiate sparse matrix operations used in Hamiltonian engineering ( Hao Xie & Wang ) , people defined 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 defined backward rules for singular value decomposition ( SVD ) function and QR decomposition ( Golub & Van Loan , 2012 ) . Instead of defining 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 differentiating 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 significantly , 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-specific ) 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 flow ( 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 flexible time-space trade-off scheme that different 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-efficient . Landauer proves that only when a device does not erase information ( i.e . reversible ) , its energy efficiency can go beyond the thermal dynamic limit . ( Landauer , 1961 ; Reeb & Wolf , 2014 ) However , these reversible programming languages can not be used directly in real scientific computing , since most of them do not have basic elements like floating point numbers , arrays , and complex numbers . This motivates us to build a new embedded domain-specific language ( eDSL ) in Julia ( Bezanson et al. , 2012 ; 2017 ) as a new playground of GP-AD . In this paper , we first compare the time-space trade-off in the optimal checkpointing and the optimal reverse computing in Sec . 2 . Then we introduce the language design of NiLang in Sec . 3 . In Sec . 4 , we explain the implementation of automatic differentiation in NiLang . In Sec . 5 , we benchmark the performance of NiLang ’ s AD with other AD software and explain why it is fast . 2 Reverse computing as an Alternative to Checkpointing One can use either checkpointing or reverse computing to trace back intermediate states of a T - step computational process s1 = f1 ( s0 ) , s2 = f2 ( s1 ) , . . . , sT = fT ( sT−1 ) with a run-time memory S . In the checkpointing scheme , the program first takes snapshots of states at certain time steps S = { sa , sb , . . . } , 1 ≤ a < b < ... ≤ T by running a forward pass . When retrieving a state sk , if sk ∈ S , just return this state , otherwise , return max j s j < k ∈ S and re-compute sk from s j . In the reverse computing scheme , one first writes the program in a reversible style . Without prior knowledge , a regular program can be transpiled to the reversible style is by doing the transformation in Listing . 1 . Listing 1 : Transpiling a regular code to the reversible code without prior knowledge . s1 += f1 ( s0 ) s2 += f2 ( s1 ) . . . sT += fT ( sT−1 ) Listing 2 : The reverse of Listing . 1 sT −= fT ( sT−1 ) . . . s2 −= f2 ( s1 ) s1 −= f1 ( s0 ) Then one can visit states in the reversed order by running the reversed program in Listing . 2 , which erases the computed results from the tail . One may argue that easing through uncomputing is not necessary here . This is not true for a general reversible program , because the intermediate states might be mutable and used in other parts of the program . It is easy to see , both checkpointing and reverse computing can trace back states without time overhead , but both suffer from a space overhead that linear to time ( Table 1 ) . The checkpointing scheme snapshots the output in every step , and the reverse computing scheme allocates extra storage for storing outputs in every step . On the other side , only checkpointing can achieve a zero space overhead by recomputing everything from the beginning s0 , with a time complexity O ( T 2 ) . The minimum space complexity in reverse computing is O ( S log ( T/S ) ) ( Bennett , 1989 ; Levine & Sherman , 1990 ; Perumalla , 2013 ) , with time complexity O ( T 1.585 ) . The difference in space overheads can be explained by the difference of the optimal checkpointing and optimal reverse computing algorithms . The optimal checkpointing algorithm that widely used in AD is the treeverse algorithm in Fig . 1 ( a ) . This algorithm partitions the computational process binomially into d sectors . At the beginning of each sector , the program snapshots the state and push it into a global stack , hence the memory for checkpointing is dS . The states in the last sector are retrieved by the above space-efficient O ( T 2 ) algorithm . After that , the last snapshot can be freed and the program has one more quota in memory . With the freed memory , the second last sector can be further partition into two sectors . Likewise , the lth sectors is partitioned into l sub-sectors , where l is the sector index counting from the tail . Recursively apply this treeverse algorithm t times until the sector size is 1 . The approximated overhead in time and space are Tc = tT , S c = dS , ( 1 ) where T = η ( t , d ) holds . By carefully choosing either a t or d , the overhead in time and space can be both logarithmic . On the other side , the optimal time-space trade-off scheme in reverse computing is the Bennett ’ s algorithm illustrated in Fig . 1 ( b ) . It evenly evenly partition the program into k sectors . The program marches forward ( P process ) for k steps to obtain the final state sk+1 , then backward ( Q process ) from the k − 1th step to erase the states in between s1 < i < k . This process is also called the compute-copyuncompute process . Recursively apply the compute-copy-uncompute process for each sector until each P/Q process contains only one unit computation . the time and space complexities are Tr = T ( T S ) ln ( 2− ( 1/k ) ) ln k , S r = k − 1 ln k S log T S . ( 2 ) Here , the overhead in time is polynomial , which is worse than the treeverse algorithm . The treeverse like partition does not apply here because the first sweep to create initial checkpoints without introducing any space overheads is not possible in reversible computing . The pseudo-code of Bennett ’ s time-space trade-off algorithm is shown in Listing . 3 . Listing 3 : The Bennett ’ s time-space trade-off scheme . The first argument { s1 , ... } 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 si+1 ← 0 allocates a zero state and add it to the state collection . Its inverse si+1 → 0 discards a zero cleared state from the collection . Its NiLang implementation is in Appendix A. bennett ( { s1 , ... } , k , i , len ) if len == 1 si+1 ← 0 fi ( si+1 , si ) else # P process that calls the forward program k steps bennett ( { s1 , ... } , k , i+len÷k * ( j-1 ) , len÷k ) for j=1,2 , ... , k # Q process that calls the backward program k-1 steps ~bennett ( { s1 , ... } , k , i+len÷k * ( j-1 ) , len÷k ) for j=k-1 , k-2 , ... ,1 The reverse computing does not show advantage from the above complexity analysis . But we argue the this analysis is from the worst case , which are very different to the practical using cases . First , τ ! δ ! is the binomial function . ( b ) Bennett ’ s time space trade-off scheme for reverse computing . ( Bennett , 1973 ; Levine & Sherman , 1990 ) P and Q are computing and uncomputing respectively . The pseudo-code is defined in Listing . 3. reverse computing can make use of the reversibility to save memory . 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 is compatible with effective codes , so that it fits 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 because it does not introduce global stack operations that harm the purity of functions . Fifth , reverse computing encourages users to think reversibly . In Appendix B.3 , we show reversible thinking can lead the user to a constant memory , constant time implementation of chained multiplication algorithms . 3 Language design NiLang is an embedded domain-specific language ( eDSL ) NiLang built on top of the host language Julia ( Bezanson et al. , 2012 ; 2017 ) . Julia is a popular language for scientific 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 define an eDSL in less than 2000 lines . Comparing with a regular reversible programming language , NiLang features array operations , rich number systems including floating-point numbers , complex numbers , fixed-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 floating-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 compiling process , grammar and operational semantics are described in Appendix G. The source code is also 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 Reversible functions and instructions Mathematically , any irreversible mapping y = f ( args ... ) can be trivially transformed to its reversible form y += f ( args ... ) or y Y= f ( args ... ) ( Y is the bit-wise XOR ) , where y is a pre-emptied variable . But in numeric computing with finite precision , this is not always true . The reversibility of arithmetic instruction is closely related to the number system . For integer and fixed 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 floating point numbers , none of the above operations are rigorously reversible . However , for convenience , we ignore the round-off errors in floating-point + and - operations and treat them on equal footing with fixed-point numbers in the following discussion . In Appendix F , 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 define a reversible multiplier in NiLang as in Listing . 4 . Listing 4 : A reversible multiplier julia > using NiLang julia > @ i function multiplier ( y ! : :Real , a : :Real , b : :Real ) y ! += a * b end julia > multiplier ( 2 , 3 , 5 ) ( 17 , 3 , 5 ) julia > ( ~multiplier ) ( 17 , 3 , 5 ) ( 2 , 3 , 5 ) Macro @ i generates two functions that are reversible to each other , multiplier and ∼multiplier , each defines a mapping R3 → R3 . The ! after a symbol is a part of the name , as a conversion to indicate the mutated variables . 3.2 Reversible 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 . 5 defines the complex valued accumulative log function . Listing 5 : Reversible complex valued log function y += log ( |x| ) + iArg ( x ) . @ i @ inline function ( : += ) ( log ) ( y ! : :Complex { T } , x : :Complex { T } ) where T n ← zero ( T ) n += abs ( x ) y ! .re += log ( n ) y ! .im += angle ( x ) n -= abs ( x ) n → zero ( T ) end Listing 6 : Compute-copy-uncompute version of Listing . 5 @ i @ inline function ( : += ) ( log ) ( y ! : :Complex { T } , x : :Complex { T } ) where T @ routine begin n ← zero ( T ) n += abs ( x ) end y ! .re += log ( n ) y ! .im += angle ( x ) ~ @ routine end 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 . 6 . 3.3 Reversible control flows One can define reversible if , for and while statements in a reversible program . Fig . 2 ( a ) shows the flow chart of executing the reversible if statement . There are two condition expressions in 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 . 7 and Listing . 8 . Listing 7 : Translating a reversible if statement ( forward ) branchkeeper = precondition if precondition branch A else branch B end assert branchkeeper == postcondition Listing 8 : Translating a reversible if statement ( backward ) branchkeeper = postcondition if postcondition ~ ( branch A ) else ~ ( branch B ) end assert branchkeeper == precondition Fig . 2 ( b ) shows the flow 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 . 9 and Listing . 10 . Listing 9 : Translating a reversible while statement ( forward ) assert postcondition == false while precondition loop body assert postcondition == true end Listing 10 : Translating a reversible while statement ( backward ) assert precondition == false while postcondition ~ ( loop body ) assert precondition == true 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 . 11 computes the Fibonacci number recursively and reversibly . Listing 11 : 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 -= n % 2 counter → 0 end out ! += 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 field for “ copying the precondition in this field as the postcondition ” . In the while statement , the postcondition is true only for the initial loop . Once code is proven correct , one can turn off 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 ) . 4 Reversible automatic differentiation 4.1 Back propagation original program O reverse program R gradient program for R comp ute reverse reverse gradient program for O 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 . Computing the gradient here is similar to forward mode automatic differentiation that computes ∂ [ multiple outputs ] ∂ [ single input ] . Inspired by the Julia package ForwardDiff ( Revels et al. , 2016 ) , we use the operator overloading technique to differentiate the program efficiently . In the backward pass , we wrap each output variable with a composite type GVar that containing an extra gradient field , and feed it into the reversed generic program . Instructions are multiple dispatched to corresponding gradient instructions that update the gradient field of GVar at the meantime of uncomputing . By reversing this gradient program , we can obtain the gradient program for the reversed program too . One can define the adjoint ( “ adjoint ” here means the program for back-propagating 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 . f : ( ~x , ~gx ) → ( ~y , ~gTx ∂~x ∂~y ) ( 3 ) f −1 : ( ~y , ~gy ) → ( ~x , ~gTy ∂~y ∂~x ) ( 4 ) It can be easily verified by applying the above two mappings consecutively , which turns out to be an identity mapping considering ∂~y ∂~x ∂~x ∂~y = 1 . The implementation details are described in Appendix C. In most languages , operator overloading is accompanied with significant 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 Hessians Combining forward mode AD and reverse mode AD is a simple yet efficient way to obtain Hessians . By wrapping the elementary type with Dual defined in package ForwardDiff and throwing it into the gradient program defined 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 significant 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 differentiable GPU kernels with NiLang . We will differentiate CUDA kernels with no more than extra 10 lines in the bundle adjustment benchmark in Sec . 5 . 5 Benchmarks 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 ForwardDiff and ReverseDiff . 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 6230 CPU @ 2.10GHz , and the GPU device is NVIDIA Titan V. For NiLang benchmarks , we have turned the reversibility check off 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 . 3 . The Tapenade data is obtained by executing the docker file provided by the original benchmark , which provides a baseline for comparison . 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 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 floating point number is S = ( 2 + d2 ) k + 2d + P , ( 5 ) S r = ( 3 + d2 + d ) k + 2log2k + P , ( 6 ) 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 ForwardDiff 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 different threads , and this requirement is satisfied 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 find more benchmarks in Appendix D , including differentiating sparse matrix dot product and obtaining Hessians in the graph embedding application . 6 Discussion In this work , we demonstrate a new approach to back propagates a program called reverse computing AD by designing a reversible eDSL NiLang . NiLang is a powerful tool to differentiate code from the source code level so that can be directly useful to machine learning researches . It can generate efficient backward rules , which is exemplified in Appendix E. It can also be used to differentiate reversible neural networks like normalizing flows ( Kobyzev et al. , 2019 ) to save memory , e.g . backpropagating NICE network ( Dinh et al. , 2014 ) with only constant space overheads . NiLang is most useful in solving large scale scientific problems memory efficiently . In Liu et al . ( 2020 ) , people solve the ground state problem of a 28 × 28 square lattice spin-glass by re-writing the quantum simulator with NiLang . There are some challenges in reverse computing AD too . • The native BLAS and convolution operations in NiLang are not optimized for the memory layout , and are too slow comparing with state of the art machine learning libraries . We need a better implementation of these functions in the reversible programming context so that it can be more useful in training traditional deep neural networks . • Although we show some examples of training neural networks on GPU , the shared-reading of a variable is not allowed . • NiLang ’ s IR does not have variable analysis . The uncomputing pass is not always necessary for the irreversible host language to deallocate memory . In many cases , the host language ’ s variable analysis can figure this out , but it is not guarantted . Another interesting issue is how to make use of reversible computing devices to save energy in machine learning . Reversible computing is not always more energy efficient than irreversible computing . In the time-space trade-off scheme in Sec . 2 , we show the time to uncompute a unit of memory is exponential to n as Qn = ( 2k − 1 ) n , and the computing energy also increases exponentially . On the other side , the amount of energy to erase a unit of memory is a constant . When ( 2k − 1 ) n > 1/ξ , erasing the memory irreversibly is more energy-efficient , where ξ is the energy ratio between a reversible operation ( an instruction or a gate ) and its irreversible counterpart . Acknowledgments 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 . References Martín Abadi , Ashish Agarwal , Paul Barham , Eugene Brevdo , Zhifeng Chen , Craig Citro , Greg S. Corrado , Andy Davis , Jeffrey Dean , Matthieu Devin , Sanjay Ghemawat , Ian Goodfellow , Andrew Harp , Geoffrey Irving , Michael Isard , Yangqing Jia , Rafal Jozefowicz , Lukasz Kaiser , Manjunath Kudlur , Josh Levenberg , Dandelion Mané , Rajat Monga , Sherry Moore , Derek Murray , Chris Olah , Mike Schuster , Jonathon Shlens , Benoit Steiner , Ilya Sutskever , Kunal Talwar , Paul Tucker , Vincent Vanhoucke , Vijay Vasudevan , Fernanda Viégas , Oriol Vinyals , Pete Warden , Martin Wattenberg , Martin Wicke , Yuan Yu , and Xiaoqiang Zheng . TensorFlow : Large-scale machine learning on heterogeneous systems , 2015 . URL https : //www.tensorflow.org/ . Software available from tensorflow.org . Jens Behrmann , David Duvenaud , and Jörn-Henrik Jacobsen . Invertible residual networks . CoRR , abs/1811.00995 , 2018 . URL http : //arxiv.org/abs/1811.00995 . Charles H. Bennett . Logical reversibility of computation . 1973 . URL https : //ieeexplore . ieee.org/abstract/document/5391327 . Charles H. Bennett . Time/space trade-offs for reversible computation . SIAM Journal on Computing , 18 ( 4 ) :766–776 , 1989. doi : 10.1137/0218053 . URL https : //doi.org/10.1137/0218053 . Ville Bergholm , Josh Izaac , Maria Schuld , Christian Gogolin , M. Sohaib Alam , Shahnawaz Ahmed , Juan Miguel Arrazola , Carsten Blank , Alain Delgado , Soran Jahangiri , Keri McKiernan , Johannes Jakob Meyer , Zeyue Niu , Antal Száva , and Nathan Killoran . Pennylane : Automatic differentiation of hybrid quantum-classical computations , 2018 . URL https : //arxiv.org/ abs/1811.04968 . Tim Besard , Christophe Foket , and Bjorn De Sutter . Effective extensible programming : Unleashing julia on gpus . CoRR , abs/1712.03112 , 2017 . URL http : //arxiv.org/abs/1712.03112 . Jeff Bezanson , Stefan Karpinski , Viral B. Shah , and Alan Edelman . Julia : A fast dynamic language for technical computing , 2012 . URL https : //arxiv.org/abs/1209.5145 . Jeff Bezanson , Alan Edelman , Stefan Karpinski , and Viral B. Shah . Julia : A fresh approach to numerical computing . SIAM Review , 59 ( 1 ) :65–98 , Jan 2017 . ISSN 1095-7200. doi : 10.1137/ 141000671 . URL http : //dx.doi.org/10.1137/141000671 . Tianqi Chen , Bing Xu , Chiyuan Zhang , and Carlos Guestrin . Training deep nets with sublinear memory cost . CoRR , abs/1604.06174 , 2016 . URL http : //arxiv.org/abs/1604.06174 . Laurent Dinh , David Krueger , and Yoshua Bengio . Nice : Non-linear independent components estimation , 2014 . URL https : //arxiv.org/abs/1410.8516 . GAEL Forget , J-M Campin , Patrick Heimbach , Christopher N Hill , Rui M Ponte , and Carl Wunsch . Ecco version 4 : An integrated framework for non-linear inverse modeling and global ocean state estimation . 2015 . URL https : //dspace.mit.edu/handle/1721.1/99660 . Michael P Frank . The r programming language and compiler . Technical report , MIT Reversible Computing Project Memo , 1997 . Michael P. Frank . Throwing computing into reverse . IEEE Spectrum , 54 ( 9 ) :32–37 , Sep 2017 . ISSN 0018-9235. doi : 10.1109/mspec.2017.8012237 . URL http : //dx.doi.org/10.1109/MSPEC . 2017.8012237 . Michael Patrick Frank and Thomas F Knight Jr. Reversibility for efficient computing . PhD thesis , 1999 . URL http : //citeseerx.ist.psu.edu/viewdoc/download ? doi=10.1.1 . 428.4962 & rep=rep1 & type=pdf . Gene H Golub and Charles F Van Loan . Matrix computations , volume 3 . JHU press , 2012 . Aidan N Gomez , Mengye Ren , Raquel Urtasun , and Roger B Grosse . The reversible residual network : Backpropagation without storing activations . In I. Guyon , U. V. Luxburg , S. Bengio , H. Wallach , R. Fergus , S. Vishwanathan , and R. Garnett ( eds . ) , Advances in Neural Information Processing Systems 30 , pp . 2214–2224 . Curran Associates , Inc. , 2017 . URL http : //papers.nips.cc/paper/ 6816-the-reversible-residual-network-backpropagation-without-storing-activations . pdf . Andreas Griewank . Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation . Optimization Methods and software , 1 ( 1 ) :35–54 , 1992 . URL https : //www.tandfonline.com/doi/abs/10.1080/10556789208805505 . Andreas Griewank and Andrea Walther . Evaluating derivatives : principles and techniques of algorithmic differentiation . SIAM , 2008 . Jin-Guo Liu Hao Xie and Lei Wang . Automatic differentiation of dominant eigensolver and its applications in quantum physics . URL https : //arxiv.org/abs/2001.04121 . Laurent Hascoet and Valérie Pascual . The tapenade automatic differentiation tool : Principles , model , and specification . ACM Transactions on Mathematical Software ( TOMS ) , 39 ( 3 ) :20 , 2013 . URL https : //dl.acm.org/citation.cfm ? id=2450158 . Tue Haulund . Design and implementation of a reversible object-oriented programming language , 2017 . URL https : //arxiv.org/abs/1707.07845 . Claudius Hubig . Use and implementation of autodifferentiation in tensor network methods with complex scalars , 2019 . URL https : //arxiv.org/abs/1907.13422 . Michael Innes . Don ’ t unroll adjoint : Differentiating ssa-form programs , 2018 . URL https : // arxiv.org/abs/1810.07951 . Michael Innes , Elliot Saba , Keno Fischer , Dhairya Gandhi , Marco Concetto Rudilosso , Neethu Mariya Joy , Tejan Karmali , Avik Pal , and Viral Shah . Fashionable modelling with flux , 2018 . URL https : //arxiv.org/abs/1811.01457 . Mike Innes , Alan Edelman , Keno Fischer , Christopher Rackauckas , Elliot Saba , Viral B. Shah , and Will Tebbutt . A differentiable programming system to bridge machine learning and scientific computing . CoRR , abs/1907.07587 , 2019 . URL http : //arxiv.org/abs/1907.07587 . Li Jing , Yichen Shen , Tena Dubcek , John Peurifoy , Scott Skirlo , Yann LeCun , Max Tegmark , and Marin Soljacić . Tunable efficient unitary neural networks ( eunn ) and their application to rnns , 2016 . URL https : //arxiv.org/abs/1612.05231 . Diederik P Kingma and Jimmy Ba . Adam : A method for stochastic optimization . URL https : //arxiv.org/abs/1412.6980 . Ivan Kobyzev , Simon Prince , and Marcus A. Brubaker . Normalizing flows : An introduction and review of current methods , 2019 . J. G. Koller and W. C. Athas . Adiabatic switching , low energy computing , and the physics of storing and erasing information . In Workshop on Physics and Computation , pp . 267–270 , Oct 1992. doi : 10.1109/PHYCMP.1992.615554 . URL https : //ieeexplore.ieee.org/document/615554 . Rolf Landauer . Irreversibility and heat generation in the computing process . IBM journal of research and development , 5 ( 3 ) :183–191 , 1961 . URL https : //ieeexplore.ieee.org/document/ 5392446 . Ivan Lanese , Naoki Nishida , Adrián Palacios , and Germán Vidal . A theory of reversibility for erlang . Journal of Logical and Algebraic Methods in Programming , 100:71–97 , Nov 2018 . ISSN 2352-2208. doi : 10.1016/j.jlamp.2018.06.004 . URL http : //dx.doi.org/10.1016/ j.jlamp.2018.06.004 . Robert Y Levine and Alan T Sherman . A note on bennett ’ s time-space tradeoff for reversible computation . SIAM Journal on Computing , 19 ( 4 ) :673–677 , 1990 . URL https : //epubs.siam . org/doi/abs/10.1137/0219046 . Hai-Jun Liao , Jin-Guo Liu , Lei Wang , and Tao Xiang . Differentiable programming tensor networks . Physical Review X , 9 ( 3 ) , Sep 2019 . ISSN 2160-3308. doi : 10.1103/physrevx.9.031041 . URL http : //dx.doi.org/10.1103/PhysRevX.9.031041 . K. Likharev . Dynamics of some single flux quantum devices : I. parametric quantron . IEEE Transactions on Magnetics , 13 ( 1 ) :242–244 , January 1977 . ISSN 1941-0069. doi : 10.1109/ TMAG.1977.1059351 . URL https : //ieeexplore.ieee.org/document/1059351 . Jin-Guo Liu , Lei Wang , and Pan Zhang . Tropical tensor network for ground states of spin glasses , 2020 . URL https : //arxiv.org/abs/2008.06888 . Xiu-Zhe Luo , Jin-Guo Liu , Pan Zhang , and Lei Wang . Yao.jl : Extensible , efficient framework for quantum algorithm design , 2019 . URL https : //arxiv.org/abs/1912.10877 . Christopher Lutz . Janus : a time-reversible language . Letter to R . Landauer. , 1986 . Matthew MacKay , Paul Vicol , Jimmy Ba , and Roger B Grosse . Reversible recurrent neural networks . In S. Bengio , H. Wallach , H. Larochelle , K. Grauman , N. CesaBianchi , and R. Garnett ( eds . ) , Advances in Neural Information Processing Systems 31 , pp . 9029–9040 . Curran Associates , Inc. , 2018 . URL http : //papers.nips.cc/paper/ 8117-reversible-recurrent-neural-networks.pdf . Dougal Maclaurin , David Duvenaud , and Ryan Adams . Gradient-based hyperparameter optimization through reversible learning . In Francis Bach and David Blei ( eds . ) , Proceedings of the 32nd International Conference on Machine Learning , volume 37 of Proceedings of Machine Learning Research , pp . 2113–2122 , Lille , France , 07–09 Jul 2015 . PMLR . URL http : //proceedings.mlr.press/v37/maclaurin15.html . Ralph C. Merkle , Robert A. Freitas , Tad Hogg , Thomas E. Moore , Matthew S. Moses , and James Ryley . Mechanical computing systems using only links and rotary joints . Journal of Mechanisms and Robotics , 10 ( 6 ) , Sep 2018 . ISSN 1942-4310. doi : 10.1115/1.4041209 . URL http : //dx . doi.org/10.1115/1.4041209 . Adam Paszke , Sam Gross , Soumith Chintala , Gregory Chanan , Edward Yang , Zachary DeVito , Zeming Lin , Alban Desmaison , Luca Antiga , and Adam Lerer . Automatic differentiation in PyTorch . In NIPS Autodiff Workshop , 2017 . URL https : //openreview.net/forum ? id= BJJsrmfCZ . Kalyan S Perumalla . Introduction to reversible computing . Chapman and Hall/CRC , 2013 . David Reeb and Michael M Wolf . An improved landauer principle with finite-size corrections . New Journal of Physics , 16 ( 10 ) :103011 , 2014 . URL https : //iopscience.iop.org/article/ 10.1088/1367-2630/16/10/103011/meta . Jarrett Revels , Miles Lubin , and Theodore Papamarkou . Forward-mode automatic differentiation in julia , 2016 . URL https : //arxiv.org/abs/1607.07892 . Chase Roberts , Ashley Milsted , Martin Ganahl , Adam Zalcman , Bruce Fontaine , Yijian Zou , Jack Hidary , Guifre Vidal , and Stefan Leichenauer . Tensornetwork : A library for physics and machine learning , 2019 . URL https : //arxiv.org/abs/1905.01330 . Matthias Seeger , Asmus Hetzel , Zhenwen Dai , Eric Meissner , and Neil D. Lawrence . Autodifferentiating linear algebra , 2017 . URL https : //arxiv.org/abs/1710.08717 . V. K. Semenov , G. V. Danilov , and D. V. Averin . Negative-inductance squid as the basic element of reversible josephson-junction circuits . IEEE Transactions on Applied Superconductivity , 13 ( 2 ) :938–943 , June 2003 . ISSN 2378-7074. doi : 10.1109/TASC.2003.814155 . URL https : //ieeexplore.ieee.org/document/1211760 . Yan Shen and Yuxing Dai . Fast automatic differentiation for large scale bundle adjustment . IEEE Access , 6:11146–11153 , 2018 . URL https : //ieeexplore.ieee.org/abstract/document/ 8307241/ . David Speyer and Bernd Sturmfels . Tropical mathematics . Mathematics Magazine , 82 ( 3 ) :163– 173 , 2009 . URL https : //www.tandfonline.com/doi/pdf/10.1080/0025570X.2009 . 11953615 . Filip Srajer , Zuzana Kukelova , and Andrew Fitzgibbon . A benchmark of selected algorithmic differentiation tools on some problems in computer vision and machine learning . Optimization Methods and Software , 33 ( 4-6 ) :889–906 , 2018 . URL https : //openreview.net/forum ? id= SMCWZLzTGDN . Jun Takahashi and Anders W. Sandvik . Valence-bond solids , vestigial order , and emergent so ( 5 ) symmetry in a two-dimensional quantum magnet , 2020 . URL https : //arxiv.org/abs/2001 . 10045 . N Takeuchi , Y Yamanashi , and N Yoshikawa . Reversible logic gate using adiabatic superconducting devices . Scientific reports , 4:6354 , 2014 . URL https : //www.nature.com/articles/ srep06354 . Naoki Takeuchi , Yuki Yamanashi , and Nobuyuki Yoshikawa . Reversibility and energy dissipation in adiabatic superconductor logic . Scientific reports , 7 ( 1 ) :1–12 , 2017 . URL https : //www . nature.com/articles/s41598-017-00089-9 . Clay S Turner . A fast binary logarithm algorithm [ dsp tips & tricks ] . IEEE Signal Processing Magazine , 27 ( 5 ) :124–140 , 2010 . URL http : //www.claysturner.com/dsp/binarylogarithm . pdf . Jean Utke , Uwe Naumann , Mike Fagan , Nathan Tallent , Michelle Strout , Patrick Heimbach , Chris Hill , and Carl Wunsch . Openad/f : A modular open-source tool for automatic differentiation of fortran codes . ACM Trans . Math . Softw. , 34 ( 4 ) , July 2008 . ISSN 0098-3500. doi : 10.1145/ 1377596.1377598 . URL https : //doi.org/10.1145/1377596.1377598 . Zhou-Quan Wan and Shi-Xin Zhang . Automatic differentiation for complex valued svd , 2019 . URL https : //arxiv.org/abs/1909.02659 . Weiqiang Zhu , Kailai Xu , Eric Darve , and Gregory C. Beroza . A general approach to seismic inversion with automatic differentiation , 2020 . URL https : //arxiv.org/abs/2003.06027 . A NiLang implementation of Bennett ’ s time-space trade-off algorithm Listing 12 : NiLang implementation of the Bennett ’ s time-space trade-off scheme . using NiLang , Test PROG_COUNTER = Ref ( 0 ) # ( 2k-1 ) ^n PEAK_MEM = Ref ( 0 ) # n * ( k-1 ) +2 @ i function bennett ( f : :AbstractVector , state : :Dict { Int , T } , k : :Int , base , len ) where T if ( len == 1 , ~ ) state [ base+1 ] ← zero ( T ) f [ base ] ( state [ base+1 ] , state [ base ] ) @ safe PROG_COUNTER [ ] += 1 @ safe ( length ( state ) > PEAK_MEM [ ] & & ( PEAK_MEM [ ] = length ( state ) ) ) else n ← 0 n += len÷k # the P process for j=1 : k bennett ( f , state , k , base+n * ( j-1 ) , n ) end # the Q process for j=k-1 : -1:1 ~bennett ( f , state , k , base+n * ( j-1 ) , n ) end n -= len÷k n → 0 end end k = 4 n = 4 N = k ^ n state = Dict ( 1= > 1.0 ) f ( x ) = x * 2.0 instructions = fill ( PlusEq ( f ) , N ) # run the program @ instr bennett ( instructions , state , k , 1 , N ) @ test state [ N+1 ] ≈ 2.0^N & & length ( state ) == 2 @ test PEAK_MEM [ ] == n * ( k-1 ) + 2 @ test PROG_COUNTER [ ] == ( 2 * k-1 ) ^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 Cases where reverse computing shows advantage B.1 Handling effective codes Reverse computing can handling effective codes with mutable structures and arrays . For example , the affine transformation can be implemented without any overhead . Listing 13 : Inplace affine transformation . @ i function i_affine ! ( y ! : :AbstractVector { T } , W : :AbstractMatrix { T } , b : :AbstractVector { T } , x : : AbstractVector { T } ) where T @ safe @ assert size ( W ) == ( length ( y ! ) , length ( x ) ) & & length ( b ) == length ( y ! ) @ invcheckoff for j=1 : size ( W , 2 ) for i=1 : size ( W , 1 ) @ inbounds y ! [ i ] += W [ i , j ] * x [ j ] end end @ invcheckoff for i=1 : size ( W , 1 ) @ inbounds y ! [ i ] += b [ i ] end end Here , the expression following the @ safe macro is an external irreversible statement . B.2 Utilizing reversibility Reverse computing can utilize reversibility to trace back states without extra memory cost . For example , we can define the unitary matrix multiplication that can be used in a type of memoryefficient recurrent neural network . ( Jing et al. , 2016 ) Listing 14 : Two level decomposition of a unitary matrix . @ i function i_umm ! ( x ! : :AbstractArray , θ ) M ← size ( x ! , 1 ) N ← size ( x ! , 2 ) k ← 0 @ safe @ assert length ( θ ) == M * ( M-1 ) /2 for l = 1 : N for j=1 : M for i=M-1 : -1 : j INC ( k ) ROT ( x ! [ i , l ] , x ! [ i+1 , l ] , θ [ k ] ) end end end k → length ( θ ) end B.3 Encourages 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 flexibility to control how to allocate memory and which number system to use . For example , to compute the power of a positive fixed-point number and an integer , one can easily write irreversible code as in Listing . 15 Listing 15 : A regular power function . function mypower ( x : :T , n : :Int ) where T y = one ( T ) for i=1 : n y * = x end return y end Listing 16 : A reversible power function . @ i function mypower ( out , x : :T , n : :Int ) where T if ( x ! = 0 , ~ ) @ routine begin ly ← one ( ULogarithmic { T } ) lx ← one ( ULogarithmic { T } ) lx * = convert ( x ) for i=1 : n ly * = x end end out += convert ( ly ) ~ @ routine end end Since the fixed-point number is not reversible under * = , naive checkpointing would require stack operations inside a loop . With reversible thinking , we can convert the fixed-point number to logarithmic numbers to utilize the reversibility of * = as shown in Listing . 16 . Here , the algorithm to convert a regular fixed-point number to a logarithmic number can be efficient . ( Turner , 2010 ) C Implementation of AD in NiLang To backpropagate the program , we first reverse the code through source code transformation and then insert the gradient code through operator overloading . If we inline all the functions in Listing . 6 , the function body would be like Listing . 17 . The automatically generated inverse program ( i.e . ( y , x ) → ( y − log ( x ) , x ) ) would be like Listing . 18 . Listing 17 : The inlined function body of Listing . 6 . @ routine begin nsq ← zero ( T ) n ← zero ( T ) nsq += x [ i ] .re ^ 2 nsq += x [ i ] .im ^ 2 n += sqrt ( nsq ) end y ! [ i ] .re += log ( n ) y ! [ i ] .im += atan ( x [ i ] .im , x [ i ] .re ) ~ @ routine Listing 18 : The inverse of Listing . 17 . @ routine begin nsq ← zero ( T ) n ← zero ( T ) nsq += x [ i ] .re ^ 2 nsq += x [ i ] .im ^ 2 n += sqrt ( nsq ) end y ! [ i ] .re -= log ( n ) y ! [ i ] .im -= atan ( x [ i ] .im , x [ i ] .re ) ~ @ routine To compute the adjoint of the computational process in Listing . 17 , one simply insert the gradient code into its inverse in Listing . 18 . The resulting inlined code is show in Listing . 19 . Listing 19 : Insert the gradient code into Listing . 18 , the original computational processes are highlighted in yellow background . @ routine begin nsq ← zero ( GVar { T , T } ) n ← zero ( GVar { T , T } ) gsqa ← zero ( T ) gsqa += x [ i ] .re.x * 2 x [ i ] .re.g -= gsqa * nsq.g gsqa -= nsq.x * 2 gsqa -= x [ i ] .re.x * 2 gsqa → zero ( T ) nsq.x += x [ i ] .re.x ^2 gsqb ← zero ( T ) gsqb += x [ i ] .im.x * 2 x [ i ] .im.g -= gsqb * nsq.g gsqb -= x [ i ] .im.x * 2 gsqb → zero ( T ) nsq.x += x [ i ] .im.x ^2 @ zeros T ra rb rta += sqrt ( nsq.x ) rb += 2 * ra nsq.g -= n.g / rb rb -= 2 * ra ra -= sqrt ( nsq.x ) ~ @ zeros T ra rb n.x += sqrt ( nsq.x ) end y ! [ i ] .re.x -= log ( n.x ) n.g += y ! [ i ] .re.g / n.x y ! [ i ] .im.x-=atan ( x [ i ] .im.x , x [ i ] .re.x ) @ zeros T xy2 jac_x jac_y xy2 += abs2 ( x [ i ] .re.x ) xy2 += abs2 ( x [ i ] .im.x ) jac_y += x [ i ] .re.x / xy2 jac_x += ( -x [ i ] .im.x ) / xy2 x [ i ] .im.g += y ! [ i ] .im.g * jac_y x [ i ] .re.g += y ! [ i ] .im.g * jac_x jac_x -= ( -x [ i ] .im.x ) / xy2 jac_y -= x [ i ] .re.x / xy2 xy2 -= abs2 ( x [ i ] .im.x ) xy2 -= abs2 ( x [ i ] .re.x ) ~ @ zeros T xy2 jac_x jac_y ~ @ routine 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 fields , 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 . Listing 20 : Time and allocation to differentiate complex valued log . julia > using NiLang , NiLang.AD , BenchmarkTools julia > @ inline function ( ir_log ) ( x : :Complex { T } ) where T log ( abs ( x ) ) + im * angle ( x ) end julia > @ btime ir_log ( x ) setup= ( x=1.0+1.2im ) ; # native code 30.097 ns ( 0 allocations : 0 bytes ) julia > @ btime ( @ instr y += log ( x ) ) setup= ( x=1.0+1.2im ; y=0.0+0.0im ) ; # reversible code 17.542 ns ( 0 allocations : 0 bytes ) julia > @ btime ( @ instr ~ ( y += log ( x ) ) ) setup= ( x=GVar ( 1.0+1.2im , 0.0+0.0im ) ; y=GVar ( 0.1+0.2im , 1.0+0.0im ) ) ; # adjoint code 25.932 ns ( 0 allocations : 0 bytes ) 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 find out that for an irreversible device , uncomputing local variables n and nsq to zeros does not affect 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 different from how it looks . The joint functions for primitive instructions ( : += ) ( sqrt ) and ( : -= ) ( sqrt ) used above can be defined as in Listing . 21 . Listing 21 : Adjoints for primitives ( : += ) ( sqrt ) and ( : -= ) ( sqrt ) . @ i @ inline function ( : -= ) ( sqrt ) ( out ! : :GVar , x : :GVar { T } ) 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 D More Benchmarks D.1 Sparse matrices We compare the call , uncall and backward propagation time used for sparse matrix dot product and matrix multiplication in Table 2 . Their reversible implementations are shown in Listing . 22 and Listing . 23 . The computing time for backward propagation is approximately 1.5-3 times the Julia ’ s native forward pass , which is close to the theoretical optimal performance . Listing 22 : Reversible sparse matrix multiplication . using SparseArrays @ i function i_dot ( r : :T , A : :SparseMatrixCSC { T } , B : :SparseMatrixCSC { T } ) where { T } m ← size ( A , 1 ) n ← size ( A , 2 ) @ invcheckoff branch_keeper ← zeros ( Bool , 2 * m ) @ safe size ( B ) == ( m , n ) || throw ( DimensionMismatch ( `` matrices must have the same dimensions '' ) ) @ invcheckoff @ inbounds for j = 1 : n ia1 ← A.colptr [ j ] ib1 ← B.colptr [ j ] ia2 ← A.colptr [ j+1 ] ib2 ← B.colptr [ j+1 ] ia ← ia1 ib ← ib1 @ inbounds for i=1 : ia2-ia1+ib2-ib1-1 ra ← A.rowval [ ia ] rb ← B.rowval [ ib ] if ( ra == rb , ~ ) r += A.nzval [ ia ] ' * B.nzval [ ib ] end # # b move - > true , a move - > false branch_keeper [ i ] Y= ia == ia2-1 || ( ib ! = ib2-1 & & ra > rb ) ra → A.rowval [ ia ] rb → B.rowval [ ib ] if ( branch_keeper [ i ] , ~ ) INC ( ib ) else INC ( ia ) end end ~ @ inbounds for i=1 : ia2-ia1+ib2-ib1-1 # # b move - > true , a move - > false branch_keeper [ i ] Y= ia == ia2-1 || ( ib ! = ib2-1 & & A.rowval [ ia ] > B.rowval [ ib ] ) if ( branch_keeper [ i ] , ~ ) INC ( ib ) else INC ( ia ) end end end @ invcheckoff branch_keeper → zeros ( Bool , 2 * m ) end Listing 23 : Reversible sparse matrix dot-product . @ i function i_mul ! ( C : :StridedVecOrMat , A : :AbstractSparseMatrix , B : :StridedVector { T } , α : : Number , β : :Number ) where T @ safe size ( A , 2 ) == size ( B , 1 ) || throw ( DimensionMismatch ( ) ) @ safe size ( A , 1 ) == size ( C , 1 ) || throw ( DimensionMismatch ( ) ) @ safe size ( B , 2 ) == size ( C , 2 ) || throw ( DimensionMismatch ( ) ) nzv ← nonzeros ( A ) rv ← rowvals ( A ) if ( β ! = 1 , ~ ) @ safe error ( `` only β = 1 is supported , got β = $ ( β ) . '' ) end # Here , we close the reversibility check inside the loop to increase performance @ invcheckoff for k = 1 : size ( C , 2 ) @ inbounds for col = 1 : size ( A , 2 ) αxj ← zero ( T ) αxj += B [ col , k ] * α for j = SparseArrays.getcolptr ( A ) [ col ] : ( SparseArrays.getcolptr ( A ) [ col + 1 ] - 1 ) C [ rv [ j ] , k ] += nzv [ j ] * αxj end αxj -= B [ col , k ] * α end end end D.2 Graph embedding problem Graph embedding can be used to find a proper representation for an order parameter ( Takahashi & Sandvik , 2020 ) in condensed matter physics . People want to find a minimum Euclidean space dimension k that a Petersen graph can embed into , that the distances between pairs of connected vertices are l1 , and the distance between pairs of disconnected vertices are l2 , where l2 > l1 . The Petersen graph is shown in Fig . 5 . Let us denote the set of connected and disconnected vertex pairs as L1 and L2 , respectively . This problem can be variationally solved with the following loss . L = Var ( dist ( L1 ) ) + Var ( dist ( L2 ) ) + exp ( relu ( dist ( L1 ) − dist ( L2 ) + 0.1 ) ) ) − 1 ( 7 ) The first 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 l2 > l1 , 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 fix two of the vertices and optimize the positions of the rest . Otherwise , the program will find 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 l2/l1 = √ 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 ForwardDiff and NiLang to obtain Hessians , it is interesting to see how much performance we can get in differentiating the graph embedding program . In Table 3 , we show the the performance of different implementations by varying the dimension k. The number of parameters is 10k . 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 ForwardDiff gives the best performance for k ≥ 3 . E Porting NiLang to Zygote 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 24 : Porting NiLang to Zygote . julia > using Zygote , NiLang , NiLang.AD , BenchmarkTools , LinearAlgebra julia > x = randn ( 1000 ) ; julia > @ benchmark norm ' ( x ) BenchmarkTools.Trial : memory estimate : 339.02 KiB allocs estimate : 8083 -- -- -- -- -- -- -- minimum time : 228.967 µs ( 0.00 % GC ) median time : 237.579 µs ( 0.00 % GC ) mean time : 277.602 µs ( 12.06 % GC ) maximum time : 5.552 ms ( 94.00 % GC ) -- -- -- -- -- -- -- samples : 10000 evals/sample : 1 julia > @ i function r_norm ( out : :T , out2 : :T , x : :AbstractArray { T } ) where T for i=1 : length ( x ) @ inbounds out2 += x [ i ] ^2 end out += sqrt ( out2 ) end julia > Zygote . @ adjoint function norm ( x : :AbstractArray { T } ) where T # compute the forward with regular norm ( might be faster ) out = norm ( x ) # compute the backward with NiLang 's norm , element type is GVar out , δy - > ( grad ( ( ~r_norm ) ( GVar ( out , δy ) , GVar ( out^2 ) , GVar ( x ) ) [ 3 ] ) , ) end julia > @ benchmark norm ' ( x ) BenchmarkTools.Trial : memory estimate : 23.69 KiB allocs estimate : 2 -- -- -- -- -- -- -- minimum time : 4.015 µs ( 0.00 % GC ) median time : 5.171 µs ( 0.00 % GC ) mean time : 6.872 µs ( 13.00 % GC ) maximum time : 380.953 µs ( 93.90 % GC ) -- -- -- -- -- -- -- samples : 10000 evals/sample : 7 We first 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 . ( 1000 × 8 + 1000 × 8 × 2 ) /1024 ≈ 23 The later one has a doubled size because GVar has an extra gradient field . F A benchmark of round-off error in leapfrog Running reversible programming with the floating pointing number system can introduce round-off errors and make the program not reversible . The quantify the effects , we use the leapfrog integrator to compute the orbitals of planets in our solar system as a benchmark . The leapfrog interations can be represented as ~ai = G m j ( ~x j − ~xi ) ‖~xi − ~x j‖3 ( 8 ) ~vi+1/2 = ~vi−1/2 + ~ai∆t ( 9 ) ~xi+1 = ~xi + ~vi+1/2∆t ( 10 ) where G is the gravitational constant and m j is the mass of jth planet , ~x , ~v and ~a are location , velocity and acceleration respectively . The first value of velocity is v1/2 = a0∆t/2 . 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 final axes of the planets and show the results in Fig . 6 . Errors are computed by comparing with the results computed with high precision floating-point numbers . One of the key steps that introduce round-off error is the computation of acceleration . If it is implemented as in Listing . 25 , the round-off error does not bring additional effect in the reversible context , hence we see overlapping lines “ ( regular ) ” and “ ( reversible ) ” in the figure . This is because , when returning a dirty ( not exactly zero cleared due to the floating-point round-off 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 . 26 , where the same variable rb is repeatedly used for compute and uncompute , the error will accumulate on this variable . In both Float64 ( double precision floating point ) and Float32 ( single precision floating point ) benchmarks , the results show a much lower precision . Hence , simulating reversible programming with floating-point numbers does not necessarily make the results less reliable if one can avoid cumulative errors in the implementation . Listing 25 : Compute the acceleration . Compute and uncompute on ancilla rc @ i function : ( += ) ( acceleration ) ( y ! : :V3 { T } , ra : :V3 { T } , rb : :V3 { T } , mb : :Real , G ) where T @ routine @ invcheckoff begin @ zeros T d anc1 anc2 anc3 anc4 rc ← zero ( V3 { T } ) d += sqdistance ( ra , rb ) anc1 += sqrt ( d ) anc2 += anc1 ^ 3 anc3 += G * mb anc4 += anc3 / anc2 rc += rb - ra end y ! += anc4 * rc ~ @ routine end Listing 26 : Compute the acceleration . Compute and uncompute on the input variable rb . @ i function : ( += ) ( acceleration ) ( y ! : :V3 { T } , ra : :V3 { T } , rb : :V3 { T } , mb : :Real , G ) where T @ routine @ invcheckoff begin @ zeros T d anc1 anc2 anc3 anc4 d += sqdistance ( ra , rb ) anc1 += sqrt ( d ) anc2 += anc1 ^ 3 anc3 += G * mb anc4 += anc3 / anc2 rb -= ra end y ! += anc4 * rb ~ @ routine # rb is not recovered rigorously ! end G Language Description G.1 Grammar A minimum definition of NiLang ’ s grammar is s a statement s multiple statements c a constant x , y , z symbols i , n integers e julia expression σP [ x 7→ v ] environemnt with x ’ s value equal to v σP [ x 7→ nothing ] environemnt with x undefined σP , e ⇓e v a Julia expression e under environment σP is interpreted as value v σP , s ⇓p σ′P the evaluation of a statement s under environment σP generates environment σ′P σP , s ⇓−1p σ′P the reverse evaluation of a statement s under environment σP generates environment σ′P Statements s : : = ∼s | e ( d∗ ) | x← e | x→ e | @ routine s ; s∗ ; ∼ @ routine | if ( e , e ) s∗ else s∗ end | while ( e , e ) s∗ end | for x = e : e : e s∗ end | begin s∗ end Data views d : : = d . x | d [ e ] | d . e | c | x Reversible functions p : : = @ i function x ( x∗ ) s * end . is the pipe operator in Julia . Here , e is a reversible function and d . e represents a bijection of d. Function arguments are data views , where a data view is a modifiable memory . It can be a variable , a field of a data view , an array element of a data view , or a bijection of a data view . G.2 Operational Semantics The following operational semantics for the forward and backward evaluation shows how a statement is evaluated and reversed . σP , e ⇓e vANCILLA σP [ x 7→ nothing ] , x→ e ⇓p σP [ x 7→ v ] σP , e ⇓e v σP [ x 7→ v ] , x← e ⇓p σP [ x 7→ nothing ] σP , x← e ⇓p σ′P ANCILLA−1 σP , x→ e⇓−1p σ′P σP , x→ e ⇓p σ′P σP , x← e⇓−1p σ′P σP , s1 ⇓p σ′P σ′P , begin s2 · · · sn end ⇓p σ′′PBLOCK σP , begin s1 · · · sn end ⇓p σ′′P σP , begin end ⇓p σP σP , sn⇓−1p σ′P σ′P , begin s1 · · · sn−1 end ⇓−1p σ′′P BLOCK−1 σP , begin s1 · · · sn end ⇓−1p σ′′P σP , begin end ⇓−1p σP σP , e1 ⇓e n1 σP , e2 ⇓e n2 σP , e3 ⇓e n3 ( n1 < = n3 ) == ( n2 > 0 ) σP [ x 7→ n1 ] , s ⇓p σ′P σ′P , for x = e1 + e2 : e2 : e3 s end ⇓p σ′′P FOR σP , for x = e1 : e2 : e3 s end ⇓p σ′′P σP , e1 ⇓e n1 σP , e2 ⇓e n2 σP , e3 ⇓e n3 ( n1 < = n3 ) ! = ( n2 > 0 ) FOR-EXIT σP , for x = e1 : e2 : e3 s end ⇓p σP σP , for x = e3 : −e2 : e1 ∼begin s end end ⇓p σ′P FOR−1 σP , for x = e1 : e2 : e3 s end ⇓−1p σ′P σP , e1 ⇓e true σP , s1 ⇓p σ′P σ′P , e2 ⇓e trueIF-T σP , if ( e1 , e2 ) s1 else s2 end ⇓p σ′P σP , e1 ⇓e f alse σP , s2 ⇓p σ′P σ′P , e2 ⇓e f alseIF-F σP , if ( e1 , e2 ) s1 else s2 end ⇓p σ′P σP , if ( e2 , e1 ) ∼begin s1 end else ∼begin s2 end end ⇓p σ′P IF−1 σP , if ( e1 , e2 ) s1 else s2 end ⇓−1p σ′P σP , e1 ⇓e true σP , e2 ⇓e f alse σP , s ⇓p σ′P σ′P , ( e1 , e2 , s ) ⇓loop σ′′PWHILE σP , while ( e1 , e2 ) s end ⇓p σ′P σP , e2 ⇓e true σP , e1 ⇓e true σP , s ⇓p σ′P σ′P , ( e1 , e2 , s ) ⇓loop σ′′PWHILE-REC σP , ( e1 , e2 , s ) ⇓loop σ′′P σP , e2 ⇓e true σP , e1 ⇓e f alse WHILE-EXIT σP , ( e1 , e2 , s ) ⇓loop σP σP , while ( e2 , e1 ) ∼begin s end end ⇓p σ′P WHILE−1 σP , while ( e1 , e2 ) s end ⇓−1p σ′P σP , s⇓−1p σ′P UNCOMPUTE σP , ∼s ⇓p σ′P σP , s1 ⇓p σ′P σ′P , begin s end ⇓p σ′′P σ′′P , s1⇓−1p σ′′′P COMPUTE-COPY-UNCOMPUTE σP , @ routine s1 ; s ; ∼ @ routine ⇓p σ′′′P σP , @ routine s1 ; ∼begin s end ∼ @ routine ⇓p σ′P COMPUTE-COPY-UNCOMPUTE−1 σP , @ routine s1 ; s ; ∼ @ routine⇓−1p σ′P σP , di ⇓get vi ∅ [ x1 7→ v1 · · · xn 7→ vn ] , ( x1 · · · xn ) = x f ( x1 · · · xn ) ⇓e σP0 [ x1 7→ v′1 · · · xn 7→ v′n ] σPi−1 , v ′ i , di ⇓set σPi CALL σP , x f ( d1 · · · dn ) ⇓p σPn σP , ( ∼x f ) ( d1 · · · dn ) ⇓p σ′P CALL−1 σP , x f ( d1 · · · dn ) ⇓−1p σ′P σP , d ⇓e vGET-VIEW σP , d ⇓get v SET-VIEW-SYM σP , v , x ⇓set σP [ x 7→ v ] z f resh σP [ z 7→ v ] , setindex ! ( d , z , x ) ⇓e v′ σP , v′ , d ⇓set σ′PSET-VIEW-ARRAY σP , v , d [ x ] ⇓set σ′P z f resh σP [ z 7→ v ] , chfield ( d , x , z ) ⇓e v′ σP , v′ , d ⇓set σ′PSET-VIEW-FIELD σP , v , d.x ⇓set σ′P z f resh σP [ z 7→ v ] , z . ( ∼e ) ⇓e v′ σP , v′ , d ⇓set σ′PSET-VIEW-BIJECTOR σP , v , d . e ⇓set σ′P Here , chfield ( x , y , z ) is a Julia function that returns an object similar to x , but with field y modified to value z , setindex ! ( x , z , y ) is a Julia function that sets the yth element of an array x to the value of z . We do not define primitive instructions like d1 = x f ( d2 · · · dn ) , ∈ { + , − , ∗ , / , Y } because these instructions are evaluated as a regular julia expression prime ( , x f ) ( d1 , d2 · · · dn ) = ( d1 x f ( d2 · · · dn ) , d2 · · · dn ) using the above CALL rule , where prime ( , x f ) is a predefined Julia function . To reverse the call , the reverse julia function ∼prime ( , x f ) should also be properly defined . On the other side , any non-primitive NiLang function defintion will generate two Julia functions x f and ∼x f so that it can be used in a recursive definition or other reversible functions . When calling a function , NiLang does not allow the input data views mappings to the same memory , i.e . shared read and write is not allowed . If the same variable is used for shared writing , the result might incorrect . If the same variable is used for both reading and writing , the program will become irreversible . For example , one can not use x -= x to erase a variable , while coding x.y -= x.z is safe . Even if a variables is used for shared reading , it can be dangerous in automatic differentiating . The share read of a variable induces a shared write of its gradient in the adjoint program . For example , y += x * x will not give the correct gradient , but its equivalent forms z ← 0 ; z += x ; y += x * z ; z -= x ; z → 0 and y += x ^2 will . G.3 Compilation The compilation of a reversible function to native Julia functions is consisted of three stages : preprocessing , reversing and translation as shown in Fig . 7 . In the preprocessing stage , the compiler pre-processes human inputs to the reversible NiLang IR . The preprocessor removes eye candies and expands shortcuts to symmetrize the code . In the left most code box in Fig . 7 , one uses @ routine < stmt > statement to record a statement , and ∼ @ routine to insert the corresponding inverse statement for uncomputing . This macro is expanded in the preprocessing stage . In the reversing stage , based on this symmetrized reversible IR , the compiler generates reversed statements . In the translation stage , the compiler translates this reversible IR as well as its inverse to native Julia code . It adds @ assignback before each function call , inserts codes for reversibility check , and handle control flows . The @ assignback macro assigns the outputs of a function to its input data views . As a final step , the compiler attaches a return statement that returns all updated input arguments . Now , the function is ready to execute on the host language .

Summary: The paper adapts reversible computing techniques to compute gradients. The techniques presented are not new though the Julia based DSL is new. The results presented are for differentiating through a GMM. It is not clear if the technique scale to a modern day neural network models and how they will integrate into current frameworks like JAX, PyTorch or TensorFlow.