Towards Distillation Guarantees under Algorithmic Alignment for Combinatorial Optimization
This work offers a rigorous theoretical foundation for distillation in structured prediction, benefiting researchers and practitioners seeking efficient models for combinatorial optimization.
The paper provides theoretical guarantees for distilling knowledge from a large source model into a smaller graph neural network for combinatorial optimization tasks, showing that distillation succeeds efficiently when the target architecture aligns with a dynamic programming algorithm and the source model satisfies the linear representation hypothesis.
Distillation transfers knowledge from a large model trained on broad data to a smaller, more efficient model suitable for deployment. In structured prediction settings, prior knowledge about the task can guide the choice of a target architecture that is algorithmically aligned with the underlying problem. Building on recent learning-theoretic analyses of decision-tree (DT) distillation (Boix-Adsera, 2024), we study when distillation succeeds for combinatorial optimization tasks. We focus on the case where the target model is a graph neural network whose architecture is aligned with a dynamic programming (DP) algorithm for the task. Assuming that the source model is sufficiently rich, formalized through the linear representation hypothesis (LRH) (Elhage et al., 2022; Park et al., 2024), we show that the distillation problem can be solved efficiently in the complexity parameters of the DP transition function, represented as a DT. Our results provide a rigorous sufficient condition for successful distillation in the flavour of algorithmic alignment.