Array & Differentiable Programming

← Back to Home

Compilers for machine learning: derive and prove what ML frameworks hard-code.

Machine learning runs on a small set of operations (multiply these arrays, differentiate this, invert that) that a compiler ought to derive rather than have a human write out by hand. My work gives these transformations a functional foundation, so they can be optimized aggressively and proved correct rather than trusted.

  • Arrays: a normal form for array programs in which partial evaluation and common subexpression elimination become one algorithm (ECOOP'24).
  • Derivatives: automatic differentiation optimized by rewrite strategies, with Timon Böhler (FTfJP'23).
  • Inverses: a verified search for inverse functions, applied to normalizing flows, with Timon Böhler and Benedict Smit (preprint).
Highlight
  • 🏅 Distinguished Paper Compiling with Arrays, ECOOP 2024
Collaborators
Timon Böhler Jannis Brugger Mattia Cerrato Cedric Derstroff Stefan Kramer Daniel Maninger Mira Mezini Viktor Pfanschilling Benedict Smit Pascal Weisenburger

Timeline

2023FTfJP @ ECOOP · Short paper

Using Rewrite Strategies for Efficient Functional Automatic Differentiation

Timon Böhler, David Richter, Mira Mezini

Automatic differentiation with dual numbers is close to the textbook definition of a derivative and easy to prove correct, but slow, and the optimizations that fix this are order-dependent. Expressing them as rewrite rules controlled by a strategy language keeps the elegance of dual numbers while stating the optimization schedule separately. On a micro-benchmark, the rewrites optimize away the nested loops of a gradient entirely (preliminary results).

Links: FTfJP 2023 · paper · arxiv

2024ECOOP

Compiling with Arrays

David Richter, Timon Böhler, Pascal Weisenburger, Mira Mezini

Typed partial evaluation is powerful but unsafe for array code: substituting a variable used twice duplicates work, and common subexpression elimination is defeated by the scopes that nested lets, loops and branches introduce. Reading arrays as the positive dual of the function type (from polarized type theory) yields AiNF, an intermediate representation whose terms are flat and maximally fissioned, so both optimizations combine into one algorithm. Termination, type preservation and maximal fission are mechanized.

A program with nested let-bindings, where y1 is bound inside z's definition and so is not in scope where y2 is defined
The equivalent program with flat let-bindings, where y1 and y2 are both in scope and their redundancy becomes visible

Fig. 1a/1b from the paper: y1 is redundant with y2, but it is not in scope at y2's definition (until the program is flattened).

Links: ECOOP 2024 · paper · artifact · github · arxiv 🏅 Distinguished Paper

2026Unpublished draft (click to expand)

Verified Inverse Function Search for Normalizing Flows

David Richter, Timon Böhler, Benedict Smit, Pascal Weisenburger, Mira Mezini

Probabilistic modeling libraries make you write each transformation and its inverse by hand, and classical program inversion demands a local invertibility that ML code lacks. Treating inversion as a search for a path that recovers the inputs by composing partial inverses (semi-inversion), with the algorithm and its soundness proof mechanized in Lean, synthesizes inverses for standard normalizing flows (additive coupling, residual, autoregressive) that existing probabilistic-programming tools fail to invert.

A table comparing SPPL, lambda-PSI, Hakaru, Stochaskell and this work across additive coupling, residual and autoregressive flows; only this work succeeds on all three
Table 1 from the paper: exact density inference tools against this algorithm. ● success, ◐ needs known weights, ○ failure.

Links: preprint

In collaboration: equation discovery

The same idea, run backwards: instead of compiling learning, use learning to search for a program (the equation behind a data set). A neural network proposes the shape of a law, and symbolic search fills in the details.

DS'25

Prompting Neural-Guided Equation Discovery Based on Residuals. Brugger, Pfanschilling, Richter, Mezini, Kramer. paper · arxiv

2025

Neural-Guided Equation Discovery. Brugger, Cerrato, Richter, Derstroff, Maninger, Mezini, Kramer. Preprint. arxiv

Open questions

Directions that come out of the work above, from thesis-sized (BSc/MSc) to PhD-scale:

  • Differentiation in AiNF. Extending the array language with automatic differentiation, proved correct, and using its optimization to remove the redundancies that differentiation generates.
  • Probabilistic primitives. Adding them to the same array language, again with correctness proofs.

Useful background is functional programming and some Lean or Scala; my Type Systems lecture covers the foundations. If one of these interests you, write me an email.