Machine learning over the p-adics

A brief introduction to non-Archimedean optimisation

Paul Lezeau (Imperial College London - LSGNT)

Joint work with Yiannis Fam, Anthea Monod and Yue Ren.

Table of contents

  1. Some background about ML & motivation

  2. Polydisc spaces

  3. Optimisation on polydisc spaces

  4. Applications

Machine learning, in brief

Usual setup in machine learning:

Data

  • Inputs: x_1, \ldots, x_n \in \RR^d

  • Outputs: y_1, \ldots, y_n \in \RR^e

For example, x_i could describe a house and y_i its price, or x_i could be a picture and y_i indicate whether it contains a kitten.

A kitten y_1 = 1

A landscape without a kitten y_2 = 0

Goal: “learn” the relation between the x_i and the y_i.

The machine-learning recipe

  • (Model) Choose a parametrised family of functions: f_\theta : \RR^d \to \RR^e Here \theta is the parameter, and we want f_\theta(x_i) \approx y_i.

  • (Data) Create a dataset of examples (\mathbf{x}, \mathbf{y}) = \{(x_i, y_i)\}_{i=1}^n.

  • (Loss) Choose a loss function \ell(\theta) = \ell(\theta; \mathbf{x}, \mathbf{y})

  • (Optimisation) Find a parameter \theta that minimises the loss.

Example: regression

  • Model: E.g. f_\theta(x) = Ax + b (or something more complicated)

  • Parameter: \theta = \{A, b\}

  • Loss: \ell(\theta; \mathbf{x}, \mathbf{y}) = \sum_{i=1}^{n} |y_i - f_\theta(x_i)|^2

Neural networks

Example. A neural network is a function of the form

f_\theta(x) = \sigma \circ L_d \circ \sigma \circ L_{d-1} \circ \cdots \circ L_1(x)

where each L_i is an affine map and \sigma is non-linear, for example

\sigma(t) = \max(0,t)

or

\sigma(t) = \frac{1}{1+e^{-t}}.

Optimisation (minimisation of the loss)

We can view loss minimisation as the gradient-flow ODE \frac{d\theta}{dt} = - \nabla_\theta \ell(\theta)

A smooth gradient-flow trajectory descending over a bumpy three-dimensional loss landscape

In practice, we discretise this to get the gradient-descent algorithm.

Gradient descent

Discrete gradient-descent iterates jumping across a bumpy three-dimensional loss landscape

Classical gradient descent

Algorithm:

  1. Pick an initial parameter \theta=\theta_0 \in \RR^n.

  2. For n=0,1,2,\ldots, repeat:

  3. Compute \nabla_\theta f at \theta_n.

  4. Set \theta_{n+1} to be \theta_n-\alpha\nabla f(\theta_n).

  5. Go back to step 2.

In practice: \nabla f is computed by backpropagation using specialised hardware (GPUs, TPUs, etc.).

If f is complicated, this can be expensive.

What about the p-adics?

Why do we care?

  • Maths
    • Useful for computations.
    • Interesting theoretical questions about algorithms.
    • ''Learning'' mathematical invariants and formulating conjectures. E.g. murmurations of elliptic curves, etc.
  • Fun example: solving systems of polynomials.
  • Applications
    • Hierarchical data is non-Archimedean
Darwin's 1837 I think sketch of an evolutionary tree
Darwin's notebook

What about the p-adics?

What changes?

  • \QQ_p is not path-connected; in fact it is totally disconnected.

  • We can have locally constant functions which are not constant.

  • We don't have a good notion of gradient.

Non-Archimedean geometry

  • (closed) points on \BA^n_K are homomorphisms K[T_1, \ldots, T_n] \to K.

  • The Berkovich affine space \BA^n_{\mathrm{Berk},K} is the set of multiplicative seminorms |\cdot| : K[T_1, \ldots, T_n] \to \RR_{\geq 0} extending the norm on K.

  • (there is a more general notion of Berkovich analytic spaces)

Example

If a \in \CC_p^n and r \in \RR_{\geq 0}^n, we can define the multiplicative seminorm | \cdot | _{B(a,r)} : \QQ_p[T_1, \ldots, T_n] \to \RR_{\geq 0} via |f|_{B(a,r)} \coloneqq \sup_{z \in B(a,r)} |f(z)|_K.

Here, B(a,r) is the closed polydisc in \CC_p^n of radius r around a.

One can check that |f|_{B(a,r)} = \sup_{\alpha \in \N^n} |c_\alpha| r^\alpha, where f = \sum_\alpha c_\alpha (T-a)^\alpha.

When r = 0 this is evaluation at a point.

Berkovich spaces

  • Berkovich spaces can be defined in general by gluing affinoid spaces.

  • These are the analogues of affine schemes in scheme theory.

  • For a variety X over \hat K, we can define the Berkovich analytification X^{\an}.

  • There are comparison theorems for the étale cohomology of X and X^{\an}.

The Berkovich affine line

Berkovich's Classification Theorem:

The discs B(a,r) \subset \CC_p, with a \in \CC_p and r \geq 0, are almost all the points of \BA^1_{\mathrm{Berk},\CC_p}.

The Berkovich affine line

Branching happens when r \in | \CC_p^\times | = p^{\QQ} (''rational points'').

Remarks

  • Over \CC_p, rational points on the affine line have infinitely many branches. This is not practical for computation.

  • The situation is much worse in higher dimensions: no such explicit description of the points of \BA^n_{\mathrm{Berk},\CC_p}.

  • In practice, we want to work with smaller fields, such as \QQ_p, and a subset \BB^n \subset \BA^n_{\mathrm{Berk},\CC_p}.

Polydisc space - Definition

Definition: The n-polydisc space is

\BB^n(\QQ_p)\coloneqq\left\{ B(a_1,r_1)\times\cdots\times B(a_n,r_n) :a_i\in\QQ_p,\ r_i\in\RR_{\geq 0}\right\}.

This has a natural topology, and several natural metrics.

If for all i, we have r_i \in p^\ZZ then say that the polydisc is rational.

Polydisc space - Definition

Example (\BB^1 for p=2)

B(·, 2¹): B(·, 2⁰): B(·, 2⁻¹): B(·, 2⁻²): B(·, 0): B(0, 2¹) B(0, 2⁰) B(0, 2⁻¹) B(1, 2⁻¹) B(0, 2⁻²) B(2, 2⁻²) B(1, 2⁻²) B(3, 2⁻²) B(0, 0) B(2, 0) B(1, 0) B(3, 0) B(3, 2⁻¹·⁵)

Polydisc space - Metric

Absolute and valuative metrics (\BB^1 for p=2)

O(p⁻¹): O(p⁰): O(p¹): O(p²): O(p∞): 0 + O(p⁻¹) 0 + O(p⁰) 0 + O(p¹) 1 + O(p¹) 0 + (p²) p + O(p²) 1 + O(p²) 1 + p + O(p²) {0} {2} {1} {3} 2⁻⁰ 2⁻¹ 2⁻² 2⁻³ + ··· 1 0.5 0.5 0.25 0.25 0.25 0.25 1 1 1 ∞ 1 1 1 1 1 1 1

Field extensions reshape disc spaces

For a valued field extension K\subseteq L, the inclusion B_K(a,r)\longmapsto B_L(a,r) embeds \BB^1(K) in \BB^1(L).

Value-group extension figure Residue-field extension figure

The value group grows
\mathbb F_2((t))\subseteq\mathbb F_2((t^{1/2}))
New intermediate vertices appear after adjoining a square root of t
New rational radii refine distances: unit edges split into half-steps.
The residue field grows
\mathbb F_2((t))\subseteq\mathbb F_4((t))
New branches appear after extending the residue field to four elements
New residue classes increase branching: degree 3 becomes degree 5.

Polydisc space - Functions

Definition (absolute polynomials)

Any polynomial f\in\QQ_p[x_1,\ldots,x_n] gives a continuous real-valued function

|f|:\BB^n\longrightarrow\RR_{\geq 0},\qquad B\longmapsto |f|_B\coloneqq\sup_{z\in B}|f(z)|_p.

Facts
  1. |f| is piecewise polynomial in the radius.
  2. |f| is monotone: |f|_B\leq |f|_{B'} whenever B\subseteq B'.
  3. a_1|f_1|+\cdots+a_r|f_r|, with a_i\in\RR, are universal approximators.
  4. These sums often have global minima (related to the projective hypersurfaces defined by the f_i).
E.g. f(T)=T^2-1 over \QQ_2.
B(1, 1/2) 1 B(1, r) −1 B(−1, r)
|T^2-1|_{B(1,r)}=\frac r2.
|T^2-1|_{B(-1,r)}=\frac r2.

Doing optimisation in practice

Julia library: NonArchimedeanMachineLearning.jl

n

p\gg 0 large

continuous optimisation

Hard!
n

p>0 small

combinatorial optimisation

MCTS, DOO, etc.

Low-exploration MCTS animation Standard-exploration MCTS animation High-exploration MCTS animation One-dimensional MCTS animation

Tree search

If we restrict to polydiscs with rational radii, minimisation becomes a tree search problem.

These are well-studied problems in computer science!

depth 0 B(0, 1) = ℤ₃ B(0, 1/3) depth 1 B(1, 1/3) B(2, 1/3) select and expand B(1, 1/9) depth 2 B(4, 1/9) B(7, 1/9) continue from here

These come in several forms:

  • deterministic

  • stochastic

Goal: explore the tree in a way that finds a minimum quickly.

Design tradeoff: need to

  • explore enough that we don't miss a minimum, but

  • avoid exploring too much and wasting time on unpromising branches.

Optimising over polydiscs is like playing chess

Computer chess: need to explore the tree of possible moves:

  • explore enough that we don't miss a good move, but

  • avoid exploring too much and wasting time on unpromising moves.

A chess game tree branching from the opening position through possible moves and replies

We can adapt one of the most successful chess algorithms, Monte Carlo tree search (MCTS).

Monte Carlo tree search, in brief

Key idea: give each node P a score (updated over time), and explore the most promising nodes first. \begin{aligned} \mathrm{score}(P) &= \text{current estimated value of }P\\ &\quad + \text{exploration bonus for }P. \end{aligned}

For each node P we associate some statistics:

  • the number of times it has been visited: N_{\text{visits}}(P)

  • the average value (or maximum) of f at its children: Q(P)

The current estimated value and the exploration bonus are functions of these.

Monte Carlo tree search

One MCTS iteration:

  • Start at the root.

  • Select the best-scoring child until we reach a leaf.

  • Evaluate f at that leaf.

  • Expand the leaf by adding its children.

  • Update scores along the path, then return to the root.

0.52 0.63 0.41 0.22 0.71 0.36 0.28 P₀ P₁ P₂ P₃ P₄ P₅ P₆ 0.52 0.63 0.71 follow the best scores evaluate f(P₄) = 0.18 ? ? P₇ P₈ update node statistics

Solving polynomials

We can use this to produce approximations to the roots of polynomials over valued fields K.

Let f \in K[x_1, \dots, x_n] be a polynomial. We want to find a \in K^n such that f(a)=0 (if it exists).

As a minimisation problem: \mathrm{argmin}_{a \in K^n} |f(a)|.

This naturally extends to the polydisc space \BB^n: \mathrm{argmin}_{P \in \BB^n} |f(P)|.

A root corresponds to the global minimum value 0.

MCTS for ''solving'' a polynomial equation

Toy examples: let f(P) = |T^2-1|_P.

MCTS

Running algorithms on this

  • Algorithms that provably converge can approximate roots of polynomials over valued fields.

  • The same approach extends to systems of polynomials and more general functions.

  • The rate of convergence appears to be related to the sizes of tubular neighbourhoods of the hypersurfaces defined by the polynomials.

What’s next?

  • Theoretical: interesting mathematical questions about algorithms.

  • Test these techniques with more real-world data.

  • Benchmark the algorithms: are non-Archimedean techniques faster than the alternatives?

  • What other classes of functions might we want to optimise?