Kotlin𝛁: Differentiable Functional Programming with Algebraic Data Types

Kotlin is a multi-platform programming language with compiler support for JVM, JS and native targets. The language emphasizes static typing, null-safety and interoperability with Java and JavaScript. In this work, we present an algebraically grounded implementation of forward and reverse mode automatic differentiation written in pure Kotlin and a property-based test suite for soundness checking. Our approach enables users to target multiple platforms through a single codebase and receive compile-time static analysis. A working prototype is provided at: https://github.com/breandan/kotlingrad

Transcript

1. Kotlin∇ Differentiable functional programming with algebraic data types Breandan Considine Universit´ e de Montr´ eal breandan.considine@umontreal.ca January 15, 2019 Breandan Considine (UdeM) Kotlin∇ January 15, 2019 1 / 16
2. Overview 1 Introduction and motivation 2 Architectural Overview 3 Usage 4 Future Plans Breandan Considine (UdeM) Kotlin∇ January 15, 2019 2 / 16
3. Type checking automatic differentiation Suppose we have a program P : R → R where: P(x) = pn ◦ pn−1 ◦ pn−2 ◦ ... ◦ p1 ◦ p0 (1) From the chain rule of calculus, we know that: dP dp0 = n i=1 dpi dpi−1 (2) In order for P to type check, what is the type of p0<i<n? pi : Tout(pi−1) → Tin(pi+1) (3) What happens if we let P : Rc → R, P : Rc → Cd or P : Ψp → Ωq? Breandan Considine (UdeM) Kotlin∇ January 15, 2019 3 / 16
4. Why Kotlin? Goal: To implement automatic differentiation in Kotlin Kotlin is a language with strong static typing and null safety Supports first-class functions, higher order functions and lambdas Has support for algebraic data types, via tuples sealed classes Extension functions, operator overloading other syntax sugar Offers features for embedding domain specific languages (DSLs) Access to all libraries and frameworks in the JVM ecosystem Multi-platform and cross-platform (JVM, Android, iOS, JS, native) Breandan Considine (UdeM) Kotlin∇ January 15, 2019 4 / 16
5. Kotlin∇ Priorities Type system Strong type system based on algebraic principles Leverage the compiler for static analysis No implicit broadcasting or shape coercion Parameterized numerical types and arbitary-precision Design principles Functional programming and lazy numerical evaluation Eager algebraic simplification of expression trees Operator overloading and tapeless reverse mode AD Usage desiderata Generalized AD with imperative array programming Automatic differentiation with infix and Polish notation Partials and higher order derivatives and gradients Testing and validation Numerical gradient checking and property-based testing Performance benchmarks and thorough regression testing Breandan Considine (UdeM) Kotlin∇ January 15, 2019 5 / 16
6. Algebraic types Abstract algebra can be useful when generalizing to new structures Helps us to easily translate between mathematics and source code Most of the time in numerical computing, we are dealing with Fields A field is a set F with two operations + and ×, with the properties: Associativity: ∀a, b, c ∈ F, a + (b + c) = (a + b) + c Commutivity: ∀a, b ∈ F, a + b = b + a and a × b = b × a Distributivity: ∀a, b, c ∈ F, a × (b × c) = (a × b) × c Identity: ∀a ∈ F, ∃0, 1 ∈ F s.t. a + 0 = a and a × 1 = a + inverse: ∀a ∈ F, ∃ − a s.t. a + (−a) = 0 × inverse: ∀a = 0 ∈ F, ∃a−1 s.t. a × a−1 = 1 Readily extensible to complex numbers, quaternions, dual numbers Field arithmetic can be implemented using parametric polymorphism What is a program, but a series of arithmetic operations? Sajovic & Vuk, Operational Calculus for Differentiable Programming Breandan Considine (UdeM) Kotlin∇ January 15, 2019 6 / 16
7. How do we define algebraic types in Kotlin∇? // T: Group<T> is effectively a self type interface Group<T: Group<T>> { operator fun plus ( f : T) : T operator fun times ( f : T) : T } // Inherits from Group, default methods interface Field <T: Field <T>>: Group<T> { operator fun unaryMinus ( ) : T operator fun minus ( f : T) : T = t h i s + −f fun i n v e r s e ( ) : T operator fun div ( f : T) : T = t h i s ∗ f . i n v e r s e () } Breandan Considine (UdeM) Kotlin∇ January 15, 2019 7 / 16
8. Algebraic Data Types class Var : Expr () class Const ( val num : Number ) : Expr () class Sum( val e1 : Expr , val e2 : Expr ) : Expr () class Prod ( val e1 : Expr , val e2 : Expr ) : Expr () sealed class Expr : Group { fun d i f f () = when( expr ) { i s Const −> Zero i s Sum −> e1 . d i f f () + e2 . d i f f () i s Prod −> e1 . d i f f () ∗ e2 + e1 ∗ e2 . d i f f () i s Var −> One } operator fun plus ( e : Expr ) = Sum( this , e ) operator fun times ( e : Expr ) = Prod ( this , e ) } Breandan Considine (UdeM) Kotlin∇ January 15, 2019 8 / 16
9. Expression simplification operator fun Expr . times ( exp : Expr ) = when { t h i s i s Const && num == 0.0 −> Const ( 0 . 0 ) t h i s i s Const && num == 1.0 −> exp exp i s Const && exp . num == 0.0 −> exp exp i s Const && exp . num == 1.0 −> t h i s t h i s i s Const && exp i s Const −> Const (num∗exp . num) else −> Prod ( this , e ) } // Sum(Prod(Const(2.0), Var()), Const(6.0)) val q = Const ( 2 . 0 ) ∗ Sum( Var () , Const ( 3 . 0 ) ) Breandan Considine (UdeM) Kotlin∇ January 15, 2019 9 / 16
10. Extension functions and contexts class Expr<T: Group<T>>: Group<Expr<T>> { //... operator fun plus ( exp : Expr<T>) = Sum( this , exp ) operator fun times ( exp : Expr<T>) = Prod ( this , exp ) } object DoubleContext { operator fun Number . times ( exp : Expr<DoubleReal >) = Const ( toDouble ( ) ) ∗ exp } // Uses ‘*‘ operator in DoubleContext fun Expr<DoubleReal >. multiplyByTwo () = with ( DoubleContext ) { 2 ∗ t h i s } Breandan Considine (UdeM) Kotlin∇ January 15, 2019 10 / 16
11. Automatic test case generation val x = v a r i a b l e ("x") val y = v a r i a b l e ("y") val z = y ∗ ( s i n ( x ∗ y ) − x ) // Function under test val dz dx = d( z ) / d( x ) // Automatic derivative val manualDx = y ∗ ( cos ( x ∗ y ) ∗ y − 1) "dz/dx should be y * (cos(x * y) * y - 1)" { a s s e r t A l l (NumGen, NumGen) { cx , cy −> // Evaluate the results at a given seed val autoEval = dz dx ( x to cx , y to cy ) val manualEval = manualDx ( x to cx , y to cy ) // Should pass if |adEval - manualEval| < eps autoEval shouldBeApproximately manualEval } } Breandan Considine (UdeM) Kotlin∇ January 15, 2019 11 / 16
12. Usage: plotting higher derivatives of nested functions with ( DoubleFunctor ) {// Use double -precision numerics val x = v a r i a b l e () // Declare an immutable variable val y = s i n ( s i n ( s i n ( x )))/ x + s i n ( x ) ∗ x + cos ( x ) + x // Lazily compute reverse -mode automatic derivatives val dy dx = d( y ) / d( x ) val d2y dx = d( dy dx ) / d( x ) val d3y dx = d( d2y dx2 ) / d( x ) val d4y dx = d( d3y dx3 ) / d( x ) val d5y dx = d( d4y dx4 ) / d( x ) p l o t ( −10..10 , dy dx , dy2 dx , d3y dx , d4y dx , d5y dx ) } Breandan Considine (UdeM) Kotlin∇ January 15, 2019 12 / 16
13. y = sin sin sin x x + x sin x + cos x + x, dy dx , d2y dx2 , d3y dx3 , d4y dx4 , d5y dx5 Breandan Considine (UdeM) Kotlin∇ January 15, 2019 13 / 16
14. Further directions to explore Theory Generalization of types to higher order functions, vector spaces Dependent types via code generation to type-check tensor dimensions General programming operators and data structures Imperative define-by-run array programming syntax Parallelization and asynchrony (cf. HogWild, YellowFin) Implementation Details Closer integration with Kotlin/Java standard library Encode additional structure, i.e. function arity into type system Vectorized optimizations for matrices with certain properties Configurable forward and backward AD modes based on dimension Automatic expression refactoring for numerical stability Primitive type specialization, i.e. FloatVector <: Vector<T>? Breandan Considine (UdeM) Kotlin∇ January 15, 2019 14 / 16
15. Learn more at: http://kg.ndan.co Breandan Considine (UdeM) Kotlin∇ January 15, 2019 15 / 16
16. Special thanks Liam Paull Michalis Famelis Alexander Nozik Hanneli Tavante Breandan Considine (UdeM) Kotlin∇ January 15, 2019 16 / 16