[Paper Review] Towards a theory of model distillation

Enric Boix-Adserà|arXiv (Cornell University)|Mar 14, 2024

Explainable Artificial Intelligence (XAI)4 citations

TL;DR

This paper introduces PAC-distillation, a formal framework for model distillation analogous to PAC-learning, enabling rigorous analysis of data and computational complexity. It proposes provably efficient algorithms to distill neural networks into interpretable decision trees using the linear representation hypothesis, showing distillation can be significantly cheaper than learning from scratch.

ABSTRACT

Distillation is the task of replacing a complicated machine learning model with a simpler model that approximates the original [BCNM06,HVD15]. Despite many practical applications, basic questions about the extent to which models can be distilled, and the runtime and amount of data needed to distill, remain largely open. To study these questions, we initiate a general theory of distillation, defining PAC-distillation in an analogous way to PAC-learning [Val84]. As applications of this theory: (1) we propose new algorithms to extract the knowledge stored in the trained weights of neural networks -- we show how to efficiently distill neural networks into succinct, explicit decision tree representations when possible by using the ``linear representation hypothesis''; and (2) we prove that distillation can be much cheaper than learning from scratch, and make progress on characterizing its complexity.

Motivation & Objective

• To formalize model distillation as a rigorous learning-theoretic problem using a PAC-style framework.
• To develop provably efficient algorithms for distilling neural networks into interpretable, explicit models such as decision trees.
• To characterize the statistical and computational complexity of distillation relative to learning from scratch.
• To explore the feasibility of distilling large models, including foundation models, into human-interpretable forms.
• To identify structural assumptions—like the linear representation hypothesis—that enable efficient distillation.

Proposed method

• Define PAC-distillation as a formal learning-theoretic problem where a model from class 𝒢 approximates a source model f∈𝒻 with high probability and low error.
• Leverage query access to the source model f and the linear representation hypothesis (LRH) to extract decision tree structure from trained neural networks.
• Design a polynomial-time algorithm that constructs an explicit decision tree representation from a neural network that implicitly computes a decision tree.
• Use reductions and web-of-reductions to relate distillation complexity to known learning-theoretic classes.
• Analyze sample complexity bounds, showing that perfect distillation requires few samples when feasible, and that agnostic distillation may require more.
• Explore extensions to distillation into logical circuits, memory-augmented models, and programs for foundation models.

Experimental results

Research questions

• RQ1 Under what conditions can a neural network be efficiently distilled into a simpler, interpretable model like a decision tree?
• RQ2 How does the sample and runtime complexity of distillation compare to that of learning from scratch?
• RQ3 Can the linear representation hypothesis enable provably efficient distillation of neural networks into explicit decision trees?
• RQ4 What is the statistical and computational complexity of agnostic distillation, and when is it feasible?
• RQ5 How can distillation be extended to foundation models like LLMs, and into what classes of models?

Key findings

• PAC-distillation formalizes distillation as a learning-theoretic problem, enabling rigorous analysis of data and runtime requirements.
• The paper presents a provably efficient algorithm to distill a neural network that implicitly computes a decision tree into an explicit decision tree representation.
• Distillation can be significantly cheaper than learning from scratch, both statistically and computationally, especially when structural assumptions like the LRH hold.
• Perfect distillation requires very few samples when feasible, indicating that distillation is statistically efficient under favorable conditions.
• Agnostic distillation may require a high number of samples, highlighting a statistical-computational gap in some settings.
• The framework suggests that distillation into expressive, interpretable classes—such as logical circuits or memory-augmented models—may be viable for foundation models like LLMs.