Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Algorithmic Guarantees for Distilling Supervised and Offline RL Datasets

Aaryan Gupta, Rishi Saket, Aravindan Raghuveer

TL;DR

This work develops a provable, training-free dataset distillation framework for both supervised regression and offline RL. For linear regression in $\mathbb{R}^d$, it shows that $\tilde{O}(d^2)$ randomly sampled regressors suffice to ensure the distilled synthetic data yields MSE losses nearly identical to those on the full training data, with a matching $\Omega(d^2)$ lower bound. It extends the approach to offline RL by distilling via the Bellman loss using random linear $Q$-value predictors; under decomposable state-action embeddings, the sample complexity improves to $\tilde{O}(d^2)$ and the optimization can be convex. Experiments validate the theory, demonstrating improved performance and data efficiency over baselines in both supervised and offline RL settings, including non-linear learners in practice.

Abstract

Given a training dataset, the goal of dataset distillation is to derive a synthetic dataset such that models trained on the latter perform as well as those trained on the training dataset. In this work, we develop and analyze an efficient dataset distillation algorithm for supervised learning, specifically regression in $\mathbb{R}^d$, based on matching the losses on the training and synthetic datasets with respect to a fixed set of randomly sampled regressors without any model training. Our first key contribution is a novel performance guarantee proving that our algorithm needs only $\tilde{O}(d^2)$ sampled regressors to derive a synthetic dataset on which the MSE loss of any bounded linear model is nearly the same as its MSE loss on the given training data. In particular, the model optimized on the synthetic data has close to minimum loss on the training data, thus performing nearly as well as the model optimized on the latter. Complementing this, we also prove a matching lower bound of $Ω(d^2)$ for the number of sampled regressors showing the tightness of our analysis. Our second contribution is to extend our algorithm to offline RL dataset distillation by matching the Bellman loss, unlike previous works which used a behavioral cloning objective. This is the first such method which leverages both, the rewards and the next state information, available in offline RL datasets, without any policy model optimization. Our algorithm generates a synthetic dataset whose Bellman loss with respect to any linear action-value predictor is close to the latter's Bellman loss on the offline RL training dataset. Therefore, a policy associated with an action-value predictor optimized on the synthetic dataset performs nearly as well as that derived from the one optimized on the training data. We conduct experiments to validate our theoretical guarantees and observe performance gains.

Algorithmic Guarantees for Distilling Supervised and Offline RL Datasets

TL;DR

This work develops a provable, training-free dataset distillation framework for both supervised regression and offline RL. For linear regression in , it shows that randomly sampled regressors suffice to ensure the distilled synthetic data yields MSE losses nearly identical to those on the full training data, with a matching lower bound. It extends the approach to offline RL by distilling via the Bellman loss using random linear -value predictors; under decomposable state-action embeddings, the sample complexity improves to and the optimization can be convex. Experiments validate the theory, demonstrating improved performance and data efficiency over baselines in both supervised and offline RL settings, including non-linear learners in practice.

Abstract

Given a training dataset, the goal of dataset distillation is to derive a synthetic dataset such that models trained on the latter perform as well as those trained on the training dataset. In this work, we develop and analyze an efficient dataset distillation algorithm for supervised learning, specifically regression in , based on matching the losses on the training and synthetic datasets with respect to a fixed set of randomly sampled regressors without any model training. Our first key contribution is a novel performance guarantee proving that our algorithm needs only sampled regressors to derive a synthetic dataset on which the MSE loss of any bounded linear model is nearly the same as its MSE loss on the given training data. In particular, the model optimized on the synthetic data has close to minimum loss on the training data, thus performing nearly as well as the model optimized on the latter. Complementing this, we also prove a matching lower bound of for the number of sampled regressors showing the tightness of our analysis. Our second contribution is to extend our algorithm to offline RL dataset distillation by matching the Bellman loss, unlike previous works which used a behavioral cloning objective. This is the first such method which leverages both, the rewards and the next state information, available in offline RL datasets, without any policy model optimization. Our algorithm generates a synthetic dataset whose Bellman loss with respect to any linear action-value predictor is close to the latter's Bellman loss on the offline RL training dataset. Therefore, a policy associated with an action-value predictor optimized on the synthetic dataset performs nearly as well as that derived from the one optimized on the training data. We conduct experiments to validate our theoretical guarantees and observe performance gains.

Paper Structure

This paper contains 36 sections, 21 theorems, 60 equations, 2 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Previous Related Work
  3. Problem Definition
  4. Our Results
  5. Overview of our Techniques
  6. Experimental Evaluations
  7. Conclusions
  8. Preliminaries for Appendix
  9. Function Classes and Pseudo-Dimension
  10. Cover over $\mathcal{B}\xspace(d, r)$
  11. Carbery-Wright anti-concentration bound
  12. Lipschitzness of feature-map $\phi$
  13. Low Rank Approximation
  14. Chernoff Bound
  15. Proof of Theorem \ref{['thm:supervised-distill-1']}
  16. ...and 21 more sections

Key Result

Theorem 4.1

Let $D^{\sf sup}_{\textnormal{train}\xspace}$ the training dataset as described above. For any $\Delta > 0$ and $\delta > 0$, let $g_1, \dots, g_k$ be iid regressors sampled from $G$ for some $k = O\left({d^2}\log({d}(B+b)/{\Delta})\log(1/\delta)\right)$. Then, with probability $1-\delta$ over the c then, for all $h \in \mathcal{F}\xspace$, $L^{\textnormal{err}\xspace}_{\textnormal{mse}\xspace}(D^

Figures (2)

  • Figure 1: MSE Test losses for Homogeneous Linear Regression on the following supervised datasets: red wine quality, white wine quality, Boston Housing.
  • Figure 2: Max Evaluation Return with two-layer neural networks on the offline RL datasets created by the following environments: Cartpole, Mountain Car, Acrobot.

Theorems & Definitions (28)

  • Theorem 4.1
  • Theorem 4.2
  • Theorem 4.3
  • Theorem 4.4
  • Theorem A.1
  • Lemma A.2
  • Theorem A.3: Theorem 8 from carberywright
  • Theorem A.4: Chernoff Bound
  • Lemma B.1
  • proof
  • ...and 18 more