[论文解读] Towards a theory of model distillation

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

Explainable Artificial Intelligence (XAI)被引用 4

一句话总结

本文提出了PAC-蒸馏（PAC-distillation），一种类比PAC-学习的模型蒸馏形式化框架，支持对数据和计算复杂度的严格分析。该框架提出可证明高效的算法，利用线性表示假设（linear representation hypothesis）将神经网络蒸馏为可解释的决策树，表明蒸馏过程相较于从零开始学习可显著降低成本。

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.

研究动机与目标

• 将模型蒸馏形式化为基于PAC风格框架的严谨学习理论问题。
• 开发可证明高效的算法，将神经网络蒸馏为可解释的显式模型（如决策树）。
• 刻画蒸馏相对于从零开始学习的统计复杂度与计算复杂度。
• 探索将大型模型（包括基础模型）蒸馏为人类可解释形式的可行性。
• 识别可实现高效蒸馏的结构性假设（如线性表示假设）。

提出的方法

• 将PAC-蒸馏定义为一种形式化的学习理论问题：从类𝒢中学习的模型以高概率和低误差逼近源模型f∈𝒻。
• 利用对源模型f的查询访问以及线性表示假设（LRH），从训练好的神经网络中提取决策树结构。
• 设计一种多项式时间算法，从隐式计算决策树的神经网络中构建显式的决策树表示。
• 通过归约与归约网络，将蒸馏复杂度与已知的学习理论类别关联。
• 分析样本复杂度边界，表明当可行时，完美蒸馏仅需少量样本；而广义蒸馏可能需要更多样本。
• 探索将蒸馏扩展至逻辑电路、内存增强模型及基础模型的程序表示。

实验结果

研究问题

• RQ1 在何种条件下，可将神经网络高效蒸馏为更简单、可解释的模型（如决策树）？
• RQ2 蒸馏的样本复杂度与运行时间复杂度与从零开始学习相比如何？
• RQ3 线性表示假设是否能实现神经网络到显式决策树的可证明高效蒸馏？
• RQ4 广义蒸馏的统计与计算复杂度如何？在何种情况下可行？
• RQ5 如何将蒸馏扩展至基础模型（如大语言模型），并可将其蒸馏为何种类别的模型？

主要发现

• PAC-蒸馏将蒸馏形式化为学习理论问题，支持对数据与运行时间需求的严格分析。
• 本文提出一种可证明高效的算法，可将隐式计算决策树的神经网络蒸馏为显式决策树表示。
• 当满足结构性假设（如LRH）时，蒸馏在统计与计算层面均显著低于从零开始学习的成本。
• 当可行时，完美蒸馏仅需极少样本，表明在有利条件下蒸馏具有统计高效性。
• 广义蒸馏可能需要大量样本，凸显了在某些场景下存在的统计-计算差距。
• 该框架表明，将蒸馏扩展至表达性强且可解释的模型类别（如逻辑电路或内存增强模型）对基础模型（如LLMs）可能是可行的。