Fair Dataset Distillation via Cross-Group Barycenter Alignment

1. Problem Statement

Dataset Distillation (DD) aims to compress a large real dataset into a small synthetic dataset while preserving the predictive performance of models trained on the original data. However, when deployed in high-stakes applications (e.g., healthcare, finance), standard DD methods often fail to maintain group fairness.

The paper identifies two critical issues:

1. Bias Amplification: Distillation processes tend to amplify existing biases in the original data. Models trained on distilled data often suffer significant performance drops for minority subgroups, leading to large fairness gaps (measured by Equalized Odds Difference, EOD).
2. Root Cause Misconception: Previous work often attributed these gaps solely to group imbalance (sample size disparities). The authors demonstrate that the problem is more complex: it arises from the interaction between group imbalance and representational separation (where different demographic groups exhibit distinct predictive patterns in the feature space).
   • Standard DD aligns synthetic data to a single aggregate target (weighted by sample counts).
   • If groups are imbalanced, the target is dominated by the majority.
   • If groups have different representations, the single target cannot simultaneously minimize the distance to all subgroups, forcing the synthetic data to "choose" a side, thereby degrading performance for others.

2. Methodology: COBRA

The authors propose COBRA (Cross-grOup BaRycenter Alignment), a framework that modifies the distillation objective to be agnostic to group imbalance.

Core Concept: Group-Agnostic Barycenter

Instead of optimizing the synthetic dataset to match a weighted average of all real samples (which favors the majority), COBRA computes a barycenter of the subgroup-specific representations.

• Input: For each class $y$, the method identifies the statistics (e.g., gradients, embeddings, or training trajectories) for each demographic subgroup $a$.
• Target Calculation: It computes a barycenter $m^*_y$ that minimizes the total distance to all subgroup statistics with uniform weights ($w_a = 1/|A|$), regardless of the subgroup's size ($\pi_{a|y}$).
  $$m^*_y = \arg \min_m \sum_{a \in A} d(\Phi_{T_{a|y}}, m)$$
  Where $\Phi$ represents the class-conditional subgroup statistics and $d$ is a discrepancy measure (e.g., squared Euclidean distance).
• Distillation Objective: The synthetic dataset is optimized to match this barycentric target $m^*_y$ rather than the vanilla weighted mixture.

Theoretical Guarantee

The paper derives an upper bound on the worst-case residual (the distance between a subgroup's true representation and the distillation target).

• Theorem 4.1: Under specific conditions, COBRA guarantees that the maximum residual across all subgroups is less than or equal to that of the vanilla (weighted) target.
• By minimizing the worst-case residual, COBRA ensures that no single subgroup is disproportionately "pushed away" from the synthetic representation, thereby reducing subgroup-specific error patterns and improving Equalized Odds.

3. Key Contributions

1. Theoretical Analysis of Bias: The authors formally characterize bias amplification in DD as a result of the interaction between group imbalance and representational separation. They prove that correcting imbalance alone is insufficient if the underlying representations differ.
2. COBRA Framework: Introduction of a principled, model-agnostic strategy that replaces the standard distillation target with a group-agnostic barycenter. This method is compatible with existing distillation techniques (Gradient Matching, Distribution Matching, Feature Alignment, Trajectory Matching).
3. Comprehensive Empirical Validation: Extensive experiments across diverse datasets (synthetic spurious correlations and real-world biased datasets like UTKFace and BFFHQ) and distillation backbones (DC, DM, CAFE, IDC, MTT).

4. Experimental Results

The paper evaluates COBRA against Vanilla DD and a Fairness-aware baseline (FairDD) using Equalized Odds Difference (EOD) and Accuracy.

• Fairness Gains: COBRA consistently achieves the lowest EOD across all benchmarks.
  • Example: On C-MNIST (Background) with Distribution Matching (DM) at IPC=50, Vanilla DD had an EOD of 100.0, while COBRA reduced it to 7.46, while simultaneously increasing accuracy from 48.8% to 96.8%.
  • On BFFHQ (real-world face dataset), COBRA reduced EOD from 63.47 to 7.87 with improved accuracy.
• Robustness to Imbalance and Separation:
  • As group imbalance (SKEW) increases, Vanilla DD's EOD rises monotonically. COBRA remains stable.
  • As representational separation (GAP) increases, Vanilla DD suffers a sharp drop in fairness. COBRA mitigates this effect effectively.
• Generalization:
  • Cross-Architecture: Distilled datasets generated by COBRA maintain fairness and accuracy when training different architectures (AlexNet, VGG, ResNet) not used during distillation.
  • Low-IPC Regime: COBRA remains effective even with very few images per class (IPC=1, 3, 5), a regime where standard DD usually fails to preserve minority information.
• Ablation Studies:
  • Noise Robustness: COBRA remains effective even when group labels are noisy (up to 50% corruption) or partially missing (using pseudo-labeling).
  • Trajectory Matching: The method successfully extends to trajectory-based distillation (MTT), proving its applicability beyond static representation matching.

5. Significance and Impact

• Paradigm Shift: The paper moves the focus of fairness in dataset distillation from "reweighting losses" (post-hoc or during optimization) to "redefining the target representation." It argues that fairness must be baked into the aggregation objective itself.
• Practical Applicability: COBRA is a drop-in replacement for the target computation in existing DD pipelines. It does not require changing the downstream training procedure or the distillation algorithm's core mechanics, only the definition of the "ideal" synthetic data.
• Addressing Real-World Bias: By demonstrating success on real-world datasets (UTKFace, BFFHQ) with complex, non-synthetic biases, the work provides a viable solution for deploying compressed models in sensitive domains without sacrificing equity.
• Theoretical Insight: The derivation of the upper bound on bias amplification provides a new theoretical lens for understanding why standard compression techniques fail in fair learning scenarios.

In conclusion, COBRA demonstrates that by aligning synthetic data to a shared, group-balanced barycenter, one can simultaneously achieve high predictive accuracy and robust group fairness, effectively decoupling the trade-off often observed in dataset distillation.