Lean Machine Learning

Current state

The library is in its early stages. Its first milestone is a general framework for stochastic sequential algorithms, used to prove regret bounds for three classical stochastic bandit algorithms. Everything is built on top of Mathlib, in particular its measure theory and probability libraries.

At a glance

Four parts: a sequential learning framework, stochastic bandits, proof tactics, and material headed for Mathlib.
Two proof automation tactics for working with Markov kernels.
A tutorial in five chapters and generated API documentation.

Where to look

Source code on GitHub
API documentation, generated from the source
Tutorials for getting started
Roadmap for what comes next

Sequential learning framework

The core of the library describes an algorithm interacting with an environment over time. An algorithm chooses an action, the environment answers with feedback, and the process repeats. This is the setting shared by bandits, online learning and reinforcement learning.

Core definitions

An algorithm is a policy given by a family of Markov kernels, mapping the history of past actions and feedback to a distribution over the next action.
An environment is described in the same style, mapping the history to a distribution over the next feedback.
The predicate for an algorithm-environment sequence states that given sequences of actions and feedback have the conditional distributions produced by a given algorithm interacting with a given environment. A variant restricts this to a finite horizon.
A canonical probability space carrying such an interaction is built for every algorithm and environment, using the Ionescu-Tulcea theorem.

Classes of environments

Oblivious environments, where the next feedback depends only on the last action rather than the whole history.
Stationary environments, where that dependence is also the same at every time step.
Bayesian environments, where an unknown parameter is drawn from a prior and determines the stationary environment.
Evaluation environments, which return the value of a measurable function of the chosen action.

Tools and generic algorithms

A martingale decomposition of the sum of rewards, and results about the means of the feedback distributions.
Densities of the law of an interaction, in both the standard and the Bayesian setting.
Random sampling, uniform sampling and round-robin algorithms, with convergence results for random sampling against an evaluation environment.

Stochastic bandits

The bandit part specializes the framework to finitely many actions and real-valued rewards, and proves regret bounds for three algorithms. These are the first end-to-end results of the library.

Definitions

The gap of an action, the best arm, and the regret of a sequence of pulls.
The Bayesian counterparts of these quantities, as random variables depending on the unknown parameter.
The array model, an alternative probability space in which an infinite array of rewards is drawn in advance, one row per action.

Algorithms and regret bounds

Three algorithms are defined in the framework, each with a proved bound on its regret. In all three, rewards are sub-Gaussian with variance proxy s2 after centering.

Explore-then-commit pulls each of the K actions m times, then commits to the one with the highest empirical mean. For a horizon n of at least K times m, its expected regret is at most the sum over actions a of gap a * (m + (n - K * m) * exp (- m * gap a ^ 2 / (4 * s2))).
Upper confidence bound pulls the action with the highest optimistic estimate of its mean. With exploration constant c, its expected regret at time n is at most the sum over actions a of 8 * c * s2 * log (n + 1) / gap a + gap a * (2 + 2 * C n), where C n stays bounded as n grows when c exceeds 2.
Thompson sampling samples an action according to its posterior probability of being optimal. For rewards lying in an interval from l to u, its Bayesian regret at time n is at most (2 * K + 1) * (u - l) + 8 * sqrt (s2 * K * n * log n).

Proof automation

Reasoning about Markov kernels involves a lot of bookkeeping. Two tactics remove some of it.

Tactics

A tactic that turns an equality of kernels into the corresponding equality of morphisms in a monoidal category, so that composition and product laws can be applied by rewriting.
A tactic that lifts an equality to a common universe level, propagating through the structure of the statement. It is extensible, so new expression shapes can be registered.

Headed for Mathlib

A substantial part of the library is general probability and measure theory that has no machine learning content, and that we intend to contribute upstream. Giving back to Mathlib is a priority of the project.

Main topics

Conditional distributions, including a predicate for having a specified conditional distribution, and conditional independence.
Independence of families indexed by infinite products.
Composition of Markov kernels, integrals against a composition, and trajectories from the Ionescu-Tulcea theorem.
Sub-Gaussian and sub-exponential moment generating functions, the concentration tools behind the regret bounds. Each comes in a plain, a conditional and a kernel version.
The Kullback-Leibler divergence of composition-products, with the chain rule in integrated form.
Assorted lemmas on measurability, absolute continuity, outer measures, finite intervals and the topology of the extended non-negative reals.

Documentation

The tutorial introduces the background needed to use the library, assuming familiarity with Lean but not with its probability library.

Tutorial chapters

Installation and setup
Probability spaces and measures
Markov kernels
Stochastic processes and martingales
Defining an algorithm

Library