Compilers for Machine Learning

Derive and prove what ML frameworks hard-code, and search for the equations that generated the data.

Machine learning runs on a small set of operations (differentiate this, invert that, multiply these arrays) that a compiler ought to be able to derive rather than have a human write out by hand. That is one half of the work here: taking the transformations deep learning frameworks hard-code and giving them a proper functional foundation, so they can be optimized aggressively and proved correct rather than trusted.

The other half runs the arrow backwards. Instead of using a compiler to make learning fast, use learning to search for a program: specifically, the equation behind a data set. A neural network proposes the shape of a law, and symbolic search fills in the details. The two halves meet in the conviction that the interesting object is the program, whether a compiler is transforming it or a search is trying to find it.

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

Using Rewrite Strategies for Efficient Functional Automatic Differentiation

Timon Böhler, David Richter, Mira Mezini

Problem. Functional automatic differentiation based on dual numbers is close to the textbook definition of a derivative, which makes it easy to prove correct (and slow). Making it fast means applying a pile of optimization rules, and those are order-dependent: one optimization enables another, so the schedule matters as much as the rules. Prior work gave the rules but no control over the order.

Idea. Express the optimizations as rewrite rules and control their application with a strategy language, keeping the elegance of dual-number differentiation while stating the optimization schedule succinctly and separately. The array language is embedded in Lean.

Run time against input size for the gradient of a vector sum: the unoptimized version and hand-written Futhark both grow steeply, while this work with rewrite rules enabled stays flat near zero — Fig. 3a from the paper: with the rewrite rules on (blue), the gradient of a vector sum stays flat where both the unoptimized version and hand-written Futhark grow.

Result. On a deliberately simple vector-sum micro-benchmark the rewrite rules optimize away the nested loops in the gradient entirely, so the optimized program stays flat where both the unoptimized one and a hand-written Futhark implementation using a dual-numbers library grow. The benchmark is meant only to show that such optimizations can in principle yield asymptotic speedups, and that application-specific optimizations can be expressed as strategies a fixed compiler pipeline would miss. Preliminary, presented as a short paper.

Links: FTfJP 2023 · paper · arxiv

2024ECOOP

Compiling with Arrays

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

Problem. Linear algebra is the substrate of neural networks, but functional languages offer lists and recursion, whereas arrays demand constant-time access and bulk operations. Some languages bridge this by representing an array as a function (which works, but leaves the idea without a logical foundation). Meanwhile typed partial evaluation is powerful yet unsafe: substituting a let-bound variable used twice duplicates code, and common subexpression elimination, which should clean that up, is defeated by the scopes that nested lets, functions and branches introduce.

Idea. Interpret functions as the ordinary negative types of polarized type theory, and arrays as the dual positive version of the function type. A positive type has a single elimination form whose computational reading is pattern matching, and just as a positive product binds two variables when matched, a positive array type binds variables with multiplicity. Booleans get the same treatment via conditionally-defined variables.

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).

Result. A new intermediate representation, A_iNF (indexed administrative normal form), normal with respect to commuting conversion for both let-bindings and for-loops, so terms come out flat and maximally fissioned (a property that opens the door to further loop optimizations). On that foundation, typed partial evaluation and common subexpression elimination combine into a single algorithm. The translation from the surface language Polara to A_iNF and its normalization are mechanized, establishing termination, type preservation, and maximal fission. Distinguished Paper at ECOOP 2024.

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

2025DS · Discovery Science

Prompting Neural-Guided Equation Discovery Based on Residuals

Jannis Brugger, Viktor Pfanschilling, David Richter, Mira Mezini, Stefan Kramer

Problem. Neural-guided equation discovery systems take a data set as a prompt and predict an equation in one shot, without extensive search. That is fast, but if the answer is not what you wanted, there is little to do short of wrestling with the system. There was no targeted way to ask for a better equation.

Idea. RED (Residuals for Equation Discovery), a post-processing step. Parse the predicted equation into a syntax tree; node-based calculation rules then give the residual for each subequation: what that subequation would have had to produce for the whole formula to fit. Feed that residual back as the target variable of a new prompt, and if the system’s answer beats the old subequation on a validation set, substitute it in.

RED applied to f(x) = x1^6 + sin(x1): the original problem goes to an equation discovery system, the residual for a subequation creates a new and simpler problem, and the improved subequation replaces the old one in the syntax tree — Fig. 1 from the paper: RED disentangling x₁⁶ + sin(x₁) into simpler subproblems.

Result. RED works with any equation discovery system, is cheap to compute, and extends easily to new mathematical operations. It helps for two distinct reasons: residuals give a fresh representation of the problem, raising the odds it resembles something the pretrained model has seen; and they factor out the part of the equation that is already correct, leaving a genuinely simpler subproblem.

Links: DS 2025 · arxiv

2025Preprint (click to expand)

Neural-Guided Equation Discovery

Jannis Brugger, Mattia Cerrato, David Richter, Cedric Derstroff, Daniel Maninger, Mira Mezini, Stefan Kramer

Problem. Equation discovery asks for the formula behind a data set, and deep learning has become an attractive way to guide that search. But the design space is large and poorly mapped: which parts of such a system actually matter, and what should one even ask of them?

Idea. Build a modular system, MGMT (Multi-Task Grammar-Guided Monte-Carlo Tree Search), so the pieces can be swapped and compared rather than argued about. A context-free grammar defines the search space, a neural network guides Monte-Carlo Tree Search state by state, constants are fitted once an equation is complete, and results feed a replay buffer that trains the network, with network and search reinforcing each other.

The MGMT architecture: a data set and grammar feed neural-guided MCTS, which exchanges guidance and state with a neural network, fits constants, stores samples in a replay buffer that supplies training data, and outputs the best equations — Fig. 1 from the paper: the MGMT architecture.

Result. Seven desirable properties of equation discovery systems, and a comparison of seven architectures (RNNs, CNNs, Transformers among them) for embedding tabular data sets, tested on contrastive learning as an auxiliary task. For nearly every combination of modules supervised learning beat reinforcement learning, and using grammar rules rather than tokens as the action space proved advantageous. Two MCTS variants (risk-seeking MCTS and AmEx-MCTS) improve the search further.

Links: arxiv

2026Unpublished draft (click to expand)

Verified Inverse Function Search for Normalizing Flows

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

Problem. Probabilistic modeling libraries make you write each transformation and its inverse by hand. Automating that is an old idea that fits modern ML badly: program inversion usually demands strict local invertibility, the more flexible approaches have no mechanized correctness proofs, and exact-inference research has concentrated on features like recursion that ML architectures avoid because they map poorly onto GPUs.

Idea. Drop the demand for local invertibility and treat inversion as search: look for a computational path that recovers the inputs from the outputs by composing partial inverses (semi-inversion). Mechanize the algorithm and its soundness proof in Lean, using intrinsically scoped and typed syntax.

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.

Result. To the authors’ knowledge the first mechanized soundness proof for a semi-inversion algorithm that copes with non-invertible and partially invertible primitives. As a case study it automatically synthesizes inverses for standard normalizing flows (additive coupling, residual and autoregressive layers), handling the non-invertible arithmetic (relu, multiplication) and data dependencies that defeat traditional inversion. Existing probabilistic-programming tools fail to invert even simple flow architectures.

Links: preprint