Kokoyi: Executable LaTeX for End-to-end Deep Learning

1 INTRODUCTION . The success of deep learning is a tale of two ends . At one end is model development , which leverages the language of mathematics to define models . At the other end is model implementation , which relies on programming languages such as Python and CUDA to unleash the power of big data and efficient computation . Despite the success , the language gap between the two ends is arguably the source of many evils , such as slow prototyping , wrong/inefficient/unreadable implementations , irreproducible results in publications , and learning barriers to name but a few . A model is often implemented and re-implemented in different , incompatible deep learning frameworks , an unnecessary fragmentation from a model point of view . The mission of KOKOYI is to close the language gap between the two ends , making math-in-codes and math-in-model consistent . KOKOYI advocates a simple principle : your model is your code . KOKOYI introduces a programming language called kokoyi-lang with a programming model elevated to mathematical equations written in LATEX , a rendering language popular among deep learning researchers ( Lamport , 1994 ) . Syntactically , kokoyi-lang is a LATEX subset that is prevalent in mathematical definitions of deep learning models , making kokoyi-lang programs renderable just like LATEX documents . Semantically , kokoyi-lang formalizes the consensus of the semantics of mathematics among deep learning researchers and practitioners . KOKOYI aims at boosting productivity at a community level : researchers code a model once instead of twice ( one in a program and another in the publication ) , kokoyi-lang is designed to be framework agnostic , the model written in kokoyi-lang in a paper is executable by all and can be ported to different frameworks . As a teaser , the following ( rendered ) line of kokoyi-lang code defines the mean squared error loss function ` for a linear regression model with weight w ∈ Rd and bias b ∈ R given a dataset D : ` ( w , b , D ) ← 1 |D| ∑ ( x , y ) ∈D ( ( w · x+ b ) − y ) 2 Given kokoyi-lang programs , KOKOYI can generate executables using various backends . In its current prototype , KOKOYI generates PyTorch code because of its widespread adoption . However , compared to PyTorch ’ s deep learning frameworks that rely on Python for usability , KOKOYI can potentially avoid Python ’ s overheads without degrading user experience by translating kokoyi-lang programs into more efficient deep learning IRs such as Relay ( Roesch et al. , 2018 ) instead of Python . The level of abstraction offered by the language of mathematics can enable various code optimizations deemed hard for existing deep learning frameworks . An example is auto-batching , which can relieve users from the burden of manually vectorizing per-sample code , an often counter-intuitive endeavor mandated solely by efficiency . Previous attempts of auto-batching like TensorFlow-Fold ( Looks et al. , 2017 ) , DyNet ( Neubig et al. , 2017 ) and JAX ( Bradbury et al. , 2018 ) heavily rely on Justin-time compilation , which are limited to statically shaped samples . By contrast , KOKOYI takes a holistic design , constructs built during compilation facilitate auto-batching during execution . As a result , KOKOYI can auto-batch complex models , such as the auto-regressive attention module in Transformer ( Vaswani et al. , 2017 ) . KOKOYI focuses on model specification . Our informal survey of sampled papers from several machine learning conferences , showing KOKOYI can impact about 15 % to 66 % of the work where model development is the focus ( Appendix A ) . As such , we do not attempt to reinvent all the wheels , compiled kokoyi-lang programs are callable Python objects and integrate with other stages ( e.g . data preparation , model tuning , and model serving ) of the deep learning pipeline ; whenever possible , KOKOYI reuses efficient existing modules . We implemented a Jupyter Notebook frontend to facilitate interactive rendering and execution of kokoyi-lang programs ( ( see HTML examples in the supplementary materials ) . There is also an interactive Web server where one can prototype a model and download both the PyTorch module and the LATEX snippet . We envision supporting compilation of kokoyi-lang code embedded in LATEX documents in the future to enable executable research papers . We implemented in kokoyi-lang several important deep learning models with various complexity , and find most of our per-sample implementations , with PyTorch as backend , achieve performance comparable to manually tensorized PyTorch implementations . At the same time , models implemented in kokoyi-lang are substantially more succinct , for instance the standard Transformer ( Vaswani et al. , 2017 ) needs only 65 lines of code , saving users from the tedious and often error-prone process of manually batching computation ( Appendix B ) . Although a large-scale user experience study is yet to be conducted , initial feedbacks from a group of deep learning researchers and data scientists indicate that compared to existing deep learning frameworks , kokoyi-lang is much closer to their mathematical formulations of deep learning models . KOKOYI will open source soon . 2 KOKOYI OVERVIEW . KOKOYI aims to be pragmatic , while adhering to the “ your model is your code ” principle . It is built around and leverages the existing Python ecosystem , yet still reserves the flexibility to revert to more efficient execution when possible . Fig . 1 ( a ) depicts the work flow when using KOKOYI . Development starts in an IDE and we currently choose the Jupyter Notebook , a popular interactive environment for live coding and data rendering . Users can write a KOKOYI program – one or more mathematical equations – in a Jupyter Notebook cell , led by the % kokoyi marker to register a hook function that renders the equation via MathJax . Programming in KOKOYI uses kokoyi-lang , a LATEX dialect that is extended with syntax familiar to DL researchers . Running the cell triggers the compilation of the equation to generate a callable python object . In this example ( Fig . 1 ( b ) ) , the returned object represents the module definition of an L-layer RNN to compute the representation of a sentence S. We hook the program execution with a proper error handler to highlight the portion of the equations that causes the error . KOKOYI translates tensors and functions defined in kokoyi-lang into tensors in underlying DL frameworks and Python functions , and makes them accessible in a Python dictionary ( named kokoyi.symbols in Jupyter notebooks ) . For example , the RNN PyTorch module calls kokoyi.symbols [ `` RNN '' ] ( Fig . 1 ( c ) ) in its forward function . KOKOYI also ports many PyTorch ’ s commonly used modules ( e.g . \ReLu ) ; the porting is mostly straightforward , with the only complexities arise from supporting auto-batching . Our current prototype enables one-click auto-generation of a Pytorch module ( Fig . 1 ( c ) ) once it is specified in kokoyi-lang . This capacity is simplistic for now and the user is required to complete the rest of the initialization ( e.g . selecting the initializer ) . It is relatively straightforward to generate more initialization code in the future . The KOKOYI compiler generates PyTorch code directly and thus the execution is just like any other PyTorch programs . Note that KOKOYI can in principle completely by-pass Python by translating the program into low-level IRs such as Relay IR ( Roesch et al. , 2018 ) and LLVM IR ( Lattner & Adve , 2004 ) to enable efficient execution . Nevertheless , we choose to focus on Python code generation mainly due to its dominating ecosystem within the deep learning community . For example , a user can easily replace the model definition in Python with KOKOYI-generated codes , while still utilizing the rich code base from the Python community for data preparation , model tuning , serving and visualization , as all the examples in the supplementary do . KOKOYI itself leverages compilers specialized for Python deep learning program such as TorchScript ( PyTorch , 2021 ) to further accelerate the generated code . The design of KOKOYI compiler is covered in Section 4 . KOKOYI performs auto-batching , where the model is written out entirely from a single sample point of view , but processes batches of them after compilation ( see Section 5 ) . 3 PROGRAMMING IN KOKOYI . Our goal is to design KOKOYI in such a way that it requires almost no learning for researchers who have written papers . In a sense , kokoyi-lang is a dialect of LATEX . KOKOYI can typeset mathematics written in kokoyi-lang as if they were written in LATEX . All mathematics in this section are typeset , parsed , and compiled by KOKOYI , thereby removing the gap between math-inpaper and math-in-code . kokoyi-lang is also more than a LATEX dialect . Mathematics written in kokoyi-lang are executable programs carrying the semantics that most members of the DL community will reasonably expect when reading them . The immediate challenge we face is that math formulations are immensely flexible . Repurposing LATEX as the programming interface must therefore allow practical and efficient code generation without losing expressiveness . In contrast to programming languages that require alphanumeric variable names , kokoyi-lang allows programmers to denote variables by math symbols typeset in LATEX . The “ real estate ” of the sub- and superscription needs particular attention . After several iterations , we decide to leave them as part of variable name and impose reasonable constraints elsewhere . For instance , we choose f ( x ; θ ) instead of fθ ( x ) for learnable functions ( i.e . modules ) , and x [ i ] ( typeset as x [ i ] ) for indexing . We also introduce common Python syntax where there is no ambiguity such as raising a power and various forms of multiplication , see supplementary for concrete examples . We now proceed to introduce kokoyi-lang ’ s key constructs : tensors , functions and modules . Tensor definition A key construct in kokoyi-lang is tensor definition , which enables one to define ( ← ) a tensor as the value of an expression , with the command \gets : [ tensor ] ← [ expression ] Expressions in tensor definitions can be applications of built-in tensor operations , such as transpose ( > ) , element-wise arithmetic ( including element-wise reductions such as sum ∑ and product ∏ ) , dot product ( · ) , and indexing ( · [ · ] ) . The expression in a tensor definition can also be a tensor constructor , which gives an entry-by-entry definition of a tensor . For example , one can define the M × N attention matrix in Transformer ( Vaswani et al. , 2017 ) as { softmax ( q > [ i ] · k [ j ] √ d ) } N j=0 M i=0 { softmax ( q > [ i ] · k [ j ] √ d ) } i j=0 M i=0 ( 1 ) where d is embedding dimension , and q and k are M × d and N × d matrices consisting of query and key vectors , respectively . The one on the left is full self-attention whereas the right is auto-regressive , i.e . only attentive to tokens behind the current one ; in a sequence-to-sequence model , the first is often used in encoder and the second in decoder . One can similarly define a multi-head attention matrix , by adding another index that ranges over [ 0 , H ) , where H is the number of heads ( see supplementary for the full Transformer model ) . One can also define a tensor in kokoyi-lang by giving a recurrence relation among its parts , such as rows and columns . For example , the hidden state matrix of a vanilla recurrent neural network ( RNN ) can be given by h [ 0≤t≤L ] ← { h0 t = 0 tanh ( Wh · h [ t−1 ] +Wx · x [ t ] + b ) otherwise ( 2 ) where x is a matrix with L rows , each being the embedding of an input token , and Wh , Wx , and h0 are learnable parameters . This is an example of KOKOYI module that we will describe next . To facilitate compiler optimization , programmers need to annotate indices where recurrence relations exist with a range ( e.g . h [ 0≤t≤L ] ) , which can depend on the dimensions of other tensors as in the case of tensor constructor . Also notice the use of \begin { cases } and \end { cases } ( rendered as the big left brace { ) in Eq . ( 2 ) for branching ( Fig . 1 ( b ) ) It is common for a recurrence relation to involve more than one tensor . For example , the hidden state and memory cell of a long short-term memory ( LSTM ) network ( Hochreiter & Schmidhuber , 1997 ) depends on the network ’ s forget , input , and output gates , which in turn depends on the previous hidden state and memory cell . This recursive mutual dependency can be expressed in kokoyi-lang as f [ 1≤t≤L ] ← σ ( Uf · h+ Vf · x [ t ] + bf ) i [ 1≤t≤L ] ← σ ( Ui · h+ Vi · x [ t ] + bi ) o [ 1≤t≤L ] ← σ ( Uo · h+ Vo · x [ t ] + bo ) c̃ [ 1≤t≤L ] ← σ ( Uc · h+ Vc · x [ t ] + bc ) c [ 1≤t≤L ] ← { c0 t = 0 f [ t ] ◦ c [ t−1 ] + i [ t ] ◦ c̃ [ t ] otherwise h [ 1≤t≤L ] ← { h0 t = 0 o [ t ] ◦ tanh ( c [ t ] ) otherwise Function and module The syntax for tensor definition applies to function definitions as well . The following is a definition of the sigmoid function : σ ( x ) ← 1 1 + e−x kokoyi-lang borrows the concept of module from popular DL frameworks such as PyTorch ( Paszke et al. , 2019 ) to enable encapsulation of learnable parameters . Conceptually , a module is a function parameterized by learnable parameters that are updated iteratively , for example , by gradient descent . Consistent with the DL practice , kokoyi-lang separates learnable parameters from parameters that the function can be applied to using a semicolon : f ( x ; θ ) . Similarly , if f uses a component module g that itself contains learnable parameters , it is written as f ( x ; g ) . The syntax of module definition resembles those of LATEX packages for pseudocode . Fig . 1 ( b ) demonstrates the major concepts of kokoyi-lang with a L-layer RNN , where the final hidden representation h [ L ] embeds the input sentence S that contains the indices of embedding of the tokens . Wx , Wh , b , h0 and the embedding table D are all trainable parameters . Note that |S| is a syntactic sugar to retrieve the first dimension of a multi-dimensional tensor , whereas the more general GetShape function returns the complete shape . { 0 } 1||S combines a tensor construction ( of a zero input ) and concatenation to right-shift the input . The rest of the codes are all self-explanatory .

Summary: This paper proposes Kokoyi, which can automatically translate mathematics into Python implementations. The proposed tool consists of kokoyi-lang, a programming language with the syntax of LATEX and the semantics of deep learning mathematics, and kokoyi-lang, a compiler and runtime supporting advanced optimizations. Kokoyi is integrated with Jupyter Notebook. To measure the flexibility of Kokoyi, the authors have implemented a variety of popular DL models and performed evaluation.