Enhancing Generalization via Sharpness-Aware Trajectory Matching for Dataset Condensation

TL;DR

This paper introduces Sharpness-Aware Trajectory Matching (SATM), a novel method for dataset condensation that improves the generalization of synthetic datasets by optimizing loss landscape sharpness and objective simultaneously, with efficient hypergradient approximation.

Contribution

The paper proposes SATM, a new approach that enhances dataset condensation by addressing generalization issues through sharpness-aware optimization and efficient hypergradient computation.

Findings

SATM outperforms existing methods on various benchmarks.

It improves generalization in both in-domain and out-of-domain settings.

The approach is easy to implement and compatible with other sharpness-aware techniques.

Abstract

Dataset condensation aims to synthesize datasets with a few representative samples that can effectively represent the original datasets. This enables efficient training and produces models with performance close to those trained on the original sets. Most existing dataset condensation methods conduct dataset learning under the bilevel (inner- and outer-loop) based optimization. However, the preceding methods perform with limited dataset generalization due to the notoriously complicated loss landscape and expensive time-space complexity of the inner-loop unrolling of bilevel optimization. These issues deteriorate when the datasets are learned via matching the trajectories of networks trained on the real and synthetic datasets with a long horizon inner-loop. To address these issues, we introduce Sharpness-Aware Trajectory Matching (SATM), which enhances the generalization capability of learned synthetic datasets by optimising the sharpness of the loss landscape and objective simultaneously. Moreover, our approach is coupled with an efficient hypergradient approximation that is mathematically well-supported and straightforward to implement along with controllable computational overhead. Empirical evaluations of SATM demonstrate its effectiveness across various applications, including in-domain benchmarks and out-of-domain settings. Moreover, its easy-to-implement properties afford flexibility, allowing it to integrate with other advanced sharpness-aware minimizers. Our code will be released.