Algorithmic stability (learning theory)

This file defines the vocabulary for algorithmic stability: datasets as tensors, replace-one and remove-one perturbations, empirical error, deterministic replace-one uniform stability (UniformStableReplace), and IID sampling of datasets.

The goal here is to provide a small, reusable vocabulary that downstream developments can reuse. The definitions follow the standard event-wise / test-point-wise inequalities from the literature, stated in a way that is easy to connect to concrete algorithms and bounds.

We keep the definitions general and avoid committing to a particular hypothesis class structure: stability is useful both for classical ERM analyses and for modern training procedures, and the right ambient structure depends on the application.

Scope and design notes

Datasets as tensors

We represent a dataset of size n as a length-n spec tensor Dataset n Z := TorchLean.Tensor Z [n].

• This integrates the learning-theory layer with TorchLean’s core, shape-indexed tensor datatype (NN.Spec.Core.Tensor.Core) and keeps the “dataset has exactly n elements” invariant enforced by the type.
• Even though the underlying tensor representation is functional (Fin n → ...), we treat datasets abstractly: what matters for stability is that we can (a) access coordinate i : Fin n and (b) perform replace-one / remove-one perturbations.
• This file holds definitions only. Concrete bounds are proved in separate files, e.g. NN.MLTheory.LearningTheory.Stability.RidgeRegression1D.

Measures and expectations

For sampling, we assume [MeasurableSpace Z] and use mathlib's ProbabilityMeasure: iid μ n is the product distribution on datasets S : Dataset n Z.

References

Algorithmic stability is a standard toolbox for generalization bounds. Classic and modern references include:

• Bousquet & Elisseeff (2002), “Stability and Generalization”.
• Shalev-Shwartz et al. (2010), “Learnability, Stability and Uniform Convergence”.
• Hardt, Recht & Singer (2016), “Train faster, generalize better: Stability of stochastic gradient descent”.
• For a textbook perspective relating stability, uniform convergence, and generalization, see Shalev-Shwartz & Ben-David (2014), “Understanding Machine Learning: From Theory to Algorithms”.

Datasets

@[reducible, inline]

abbrev NN.MLTheory.LearningTheory.Stability.Dataset (n : ℕ) (Z : Type) : Type

A dataset of size n with examples in Z.

The leading dimension records the sample count in the type.

@[reducible, inline]

abbrev NN.MLTheory.LearningTheory.Stability.Dataset.toFn {n : ℕ} {Z : Type} (S : Dataset n Z) : Fin n → Z

View a dataset tensor as a function Fin n → Z.

This is definitional content via TorchLean.Tensor.vectorEquiv, and is used to:

• define replace/remove operations via Function.update and Fin.succAbove, and
• transport the standard product measurable space / IID sampling measure to the tensor type.

@[reducible, inline]

abbrev NN.MLTheory.LearningTheory.Stability.Dataset.ofFn {n : ℕ} {Z : Type} (f : Fin n → Z) : Dataset n Z

Build a dataset tensor from a function Fin n → Z.

@[simp]

theorem NN.MLTheory.LearningTheory.Stability.Dataset.toFn_ofFn {n : ℕ} {Z : Type} (f : Fin n → Z) : (ofFn f).toFn = f

Reading back a dataset built from a function recovers the function.

@[simp]

theorem NN.MLTheory.LearningTheory.Stability.Dataset.ofFn_toFn {n : ℕ} {Z : Type} (S : Dataset n Z) : ofFn S.toFn = S

The other round trip. Together with toFn_ofFn this is what lets stability arguments move freely between the tensor representation of a sample and the function view the measure theory prefers.

@[reducible, inline]

abbrev NN.MLTheory.LearningTheory.Stability.Dataset.get {n : ℕ} {Z : Type} (S : Dataset n Z) (i : Fin n) : Z

Coordinate access for dataset tensors.

@[simp]

theorem NN.MLTheory.LearningTheory.Stability.Dataset.get_ofFn {n : ℕ} {Z : Type} (f : Fin n → Z) (i : Fin n) : (ofFn f).get i = f i

Coordinate access on a dataset built from a function is just application.

@[instance_reducible]

instance NN.MLTheory.LearningTheory.Stability.Dataset.instMeasurableSpace {n : ℕ} {Z : Type} [MeasurableSpace Z] : MeasurableSpace (Dataset n Z)

The measurable space on dataset tensors is the one transported from the standard product measurable space on functions Fin n → Z.

This makes IID sampling (iid below) and the standard stability definitions work without changing their measure-theoretic content; it is just a representation choice.

theorem NN.MLTheory.LearningTheory.Stability.Dataset.measurable_toFn {n : ℕ} {Z : Type} [MeasurableSpace Z] : Measurable toFn

toFn is measurable by construction: the measurable space on datasets is its comap.

theorem NN.MLTheory.LearningTheory.Stability.Dataset.measurable_ofFn {n : ℕ} {Z : Type} [MeasurableSpace Z] : Measurable ofFn

ofFn is measurable too, so the representation choice is invisible to the measure theory.

def NN.MLTheory.LearningTheory.Stability.replaceAt {Z : Type} {n : ℕ} (S : Dataset n Z) (i : Fin n) (z' : Z) : Dataset n Z

Replace the example at index i with z'.

This is the standard “replace-one” perturbation used in uniform stability definitions.

@[simp]

theorem NN.MLTheory.LearningTheory.Stability.get_replaceAt {Z : Type} {n : ℕ} (S : Dataset n Z) (i j : Fin n) (z' : Z) : (replaceAt S i z').get j = if j = i then z' else S.get j

Reading a replaced dataset returns the replacement at that coordinate and the original example everywhere else.

def NN.MLTheory.LearningTheory.Stability.removeAt {Z : Type} {n : ℕ} (S : Dataset (n + 1) Z) (i : Fin (n + 1)) : Dataset n Z

Remove the example at index i from a dataset of size n+1.

This uses Fin.succAbove to reindex the remaining elements into Fin n.

@[simp]

theorem NN.MLTheory.LearningTheory.Stability.get_removeAt {Z : Type} {n : ℕ} (S : Dataset (n + 1) Z) (i : Fin (n + 1)) (j : Fin n) : (removeAt S i).get j = S.get (i.succAbove j)

Reading a shortened dataset skips the removed index, which is what Fin.succAbove encodes.

Learning algorithms and loss

@[reducible, inline]

abbrev NN.MLTheory.LearningTheory.Stability.LearningMap (n : ℕ) (Z H : Type) : Type

A deterministic learning algorithm mapping datasets to hypotheses.

This is the interface needed to state stability: an “algorithm” is just a function Dataset n Z → H.

@[reducible, inline]

abbrev NN.MLTheory.LearningTheory.Stability.Loss (H Z : Type) : Type

A real-valued loss function.

We fix the codomain to ℝ to match the standard stability literature.

Errors

noncomputable def NN.MLTheory.LearningTheory.Stability.empiricalError {Z H : Type} {n : ℕ} (ℓ : Loss H Z) (h : H) (S : Dataset n Z) : ℝ

Empirical error (average loss on a dataset).

We write this with an explicit $1/n$ normalization so downstream lemmas can control constants. At n = 0, the totalized real expression is zero.

Deterministic replace-one stability

def NN.MLTheory.LearningTheory.Stability.UniformStableReplace {Z H : Type} {n : ℕ} (A : LearningMap n Z H) (ℓ : Loss H Z) (β : ℝ) : Prop

Deterministic replace-one uniform stability (a common core notion).

UniformStableReplace A ℓ β means that if you replace one example in the training set, then the loss on any test point changes by at most $\beta$.

IID sampling helper

noncomputable def NN.MLTheory.LearningTheory.Stability.iid {Z : Type} [MeasurableSpace Z] (μ : MeasureTheory.ProbabilityMeasure Z) (n : ℕ) : MeasureTheory.ProbabilityMeasure (Dataset n Z)

IID sampling: product distribution on datasets.

If μ is a distribution over examples Z, then iid μ n is the distribution over datasets of size n obtained by sampling each coordinate independently from μ.

The product distribution is naturally defined on functions Fin n → Z; Dataset.ofFn transports it to the shape-indexed tensor representation.