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: The paper presents a LaTeX-based language and compiler called kokoyi to write math-based models and compile them to actual code (such as PyTorch). The authors present an approach to support optimizations such as auto-batching during this compilation process which significantly reduces user burden. The authors presented kokoyi implementations of several popular models such MLP, CNNs, LSTMs and transformers and showed that the kokoyi compiler does not introduce much performance drop.