Geometric Characterisation and Structured Trajectory Surrogates for Clinical Dataset Condensation

1. Problem Statement

Dataset Condensation aims to create compact synthetic datasets that retain the training utility of large real-world datasets, facilitating efficient model development and data sharing in governed domains like healthcare. Trajectory Matching (TM) is a leading condensation approach where synthetic data is optimized to reproduce the parameter updates (trajectories) observed when training a model on real data.

However, the paper identifies a fundamental, unresolved issue: The representational bottleneck.

• The Mismatch: Standard TM uses raw Stochastic Gradient Descent (SGD) trajectories as supervision. These trajectories often exhibit "spectrally broad" structures, meaning the parameter changes occur in many weakly aligned directions.
• The Constraint: A fixed, small synthetic dataset can only generate gradients within a limited "reachable span" (a low-dimensional subspace).
• The Consequence: When the supervision signal (the real SGD trajectory) is spectrally broad, a significant portion of the signal lies outside the synthetic dataset's reachable span. This creates an irreducible matching error, particularly problematic in low-prevalence clinical settings (e.g., rare diseases) where minority-class updates are sparse and inconsistent.

2. Methodology: Bézier Trajectory Matching (BTM)

The authors propose Bézier Trajectory Matching (BTM) to address the spectral mismatch by replacing raw SGD trajectories with structured, low-rank surrogates.

A. Geometric Characterization

The paper first provides a theoretical proof showing that the trajectory matching loss is lower-bounded by the projection error of the teacher's parameter displacement onto the student's reachable gradient span.

• Theorem: If the teacher's displacement matrix $A$ has a rank or spectral distribution that exceeds the dimension of the synthetic dataset's reachable span, a non-zero error is unavoidable.
• Implication: Effective condensation requires structuring the supervision signal to be concentrated in a few coherent directions rather than reproducing the full stochastic noise of SGD.

B. The BTM Algorithm

Instead of storing long sequences of SGD checkpoints, BTM learns quadratic Bézier curves to connect the initial model state ($\theta_0$) and the final trained state ($\theta_T$).

1. Surrogate Construction: For each training run, a quadratic Bézier curve is defined:
   $$\Phi(t) = (1-t)^2\theta_0 + 2t(1-t)\phi + t^2\theta_T$$
   where $\phi$ is a learnable control point.
2. Optimization of Control Points: The control point $\phi$ is optimized to minimize the average training loss along the path. This encourages a smooth, monotonic descent profile that avoids high-loss excursions common in raw SGD.
3. Low-Rank Supervision:
   • A single quadratic Bézier curve induces displacements that lie in a 2-dimensional subspace (spanned by $\theta_T - \theta_0$ and the control point vector).
   • This drastically reduces the effective rank of the supervision signal compared to raw SGD, aligning it with the limited representational capacity of small synthetic datasets.
4. Storage Efficiency: Storing a trajectory requires only 3 parameter vectors (start, control, end) per run, compared to $T$ checkpoints for raw SGD, reducing storage by orders of magnitude.

3. Key Contributions

1. Geometric Characterization of TM: The authors formalize trajectory matching as a constrained subspace-approximation problem, proving that a "conditional representability bottleneck" exists when the supervision signal's spectral energy exceeds the synthetic dataset's reachable span.
2. Empirical Analysis of Supervision: They demonstrate that standard SGD supervision is spectrally broad, especially in low-prevalence settings, explaining why TM often fails in rare-event clinical prediction.
3. Bézier Trajectory Matching (BTM): A novel method replacing raw SGD with task-optimized quadratic Bézier surrogates. This yields a structured, low-rank supervision signal that is better aligned with the constraints of synthetic data.
4. Comprehensive Clinical Evaluation: Extensive experiments on five diverse clinical datasets (tabular and time-series) demonstrating superior performance, particularly in low-data and low-prevalence regimes.

4. Experimental Results

The authors evaluated BTM on five clinical datasets: Oxford, Portsmouth, and Birmingham NHS cohorts (COVID-19 prediction), eICU, and MIMIC-III (mortality prediction and phenotyping).

• Performance Gains:
  • BTM consistently matches or outperforms state-of-the-art baselines (MTT, FTD, DATM, MCT).
  • Low-Prevalence Regimes: The largest gains are observed in low-prevalence settings (e.g., Birmingham dataset with 0.8% prevalence). BTM improved AUPRC by up to 17.9% over the strongest baseline in these sparse data scenarios.
  • Multi-label Tasks: On MIMIC-III phenotyping (25 classes), BTM achieved the best Macro AUPRC across all budgets, indicating superior retention of label-specific signals.
• Cross-Architecture Generalization:
  • Synthetic datasets condensed using BTM transferred more robustly to unseen model architectures (e.g., from MLP to TCN/LSTM) compared to baselines, suggesting the learned supervision captures task-relevant structure rather than architecture-specific optimization noise.
• Storage Efficiency:
  • BTM reduced trajectory storage requirements by 20x to 33x compared to storing full SGD trajectories (e.g., reducing 1340 MB to 41 MB on the Oxford dataset).
• Ablation Studies:
  • Replacing raw SGD with simple linear interpolation helped, but learned quadratic Bézier curves provided the best performance.
  • The improvement stems from the Bézier curve's ability to minimize average path loss, keeping the supervision signal in low-loss regions, unlike convexified or linear paths which may traverse high-loss areas.

5. Significance and Impact

• Theoretical Insight: The paper shifts the focus of dataset condensation from "reproducing stochastic paths" to "structuring the supervision signal." It establishes that the geometry of the gradient span is a limiting factor in condensation.
• Clinical Utility: By addressing the spectral mismatch, BTM makes dataset condensation viable for rare disease prediction and resource-constrained environments where data is scarce and privacy-governed.
• Practical Deployment: The massive reduction in storage requirements (33x) makes it feasible to store and share multiple expert trajectories, facilitating collaboration in clinical research without compromising data governance.
• Future Directions: The authors note that while BTM improves utility, it does not provide formal privacy guarantees. Future work should integrate differential privacy and bias mitigation to ensure safe deployment in healthcare.

In summary, this paper argues that effective dataset condensation depends on the structure of the supervision signal, not just its magnitude. By replacing noisy, high-rank SGD trajectories with smooth, low-rank Bézier surrogates, BTM overcomes the geometric bottlenecks of previous methods, delivering state-of-the-art results in challenging clinical prediction tasks.