Technical Summary: DriftXpress

Problem Statement

Drifting Models represent a paradigm shift in generative modeling by enabling one-step image generation without the iterative inference required by diffusion or flow-based models. Instead of iterative denoising, Drifting Models learn an attraction–repulsion field in feature space that pulls generated samples toward the data distribution and pushes them away from the current model distribution.

While effective, standard Drifting Models introduce a significant computational trade-off: inference is cheap (single-step), but training is expensive. The primary bottleneck lies in the repeated estimation of the attraction field, which requires computing exact kernel interactions between every generated sample and the entire training support (or large mini-batches) at every training step. This results in high wall-clock training costs, limiting the scalability of the approach. The paper asks whether this training-inference trade-off can be improved to reduce training costs without sacrificing the one-step inference advantage.

Methodology: DriftXpress

The authors propose DriftXpress, an accelerated formulation of drifting models based on Projected Reproducing Kernel Hilbert Space (RKHS) fields. The core idea is to approximate the attraction component of the drifting field using a low-rank feature space derived from the Nyström method, while keeping the repulsion component exact.

1. RKHS Projection via Nyström Landmarks

The method leverages the Nyström approximation to project the kernel interactions onto a subspace spanned by a fixed set of "landmark" points selected from the training data prior to training.

• Landmark Selection: A set of landmarks $U = \{u_1, \dots, u_r\}$ is sampled (e.g., uniformly per class) from the training data. These landmarks define a subspace $H_U \subset H_k$.
• Feature Map: A Nyström feature map $\phi(z)$ is constructed using the landmark kernel matrix $K_{UU}$ and the kernel evaluations between a point $z$ and the landmarks. This allows the kernel similarity $k(x, y)$ to be approximated by a low-rank dot product $\phi(x)^\top \phi(y)$.

2. Projected Attraction with Cached Summaries

In standard drifting, the attractive field $V_p^+(x)$ is computed as a normalized kernel-weighted mean of training samples. DriftXpress replaces the exact kernel $k$ with the projected kernel $k_U$.

• Precomputation: Instead of recomputing interactions with the full training set at every step, DriftXpress precomputes two compact summaries over the entire training dataset:
  • $A_p = \sum \phi(y_j^+) (y_j^+)^\top$
  • $b_p = \sum \phi(y_j^+)$
• Efficient Evaluation: During training, the attractive mean $\mu_p^U(x)$ is computed efficiently using these cached summaries and the query's projected features $\phi(x)$:
  $$\mu_p^U(x) = \frac{A_p^\top \phi(x)}{\phi(x)^\top b_p + \epsilon}$$
  This reduces the per-step complexity from depending on the training set size $N_+$ to depending on the landmark dimension $r$ (where $r \ll N_+$).

3. Exact Repulsion

The repulsive component $V_q^-(x)$, which pushes samples away from the current generated batch, is kept exact. The authors observe that approximating repulsion leads to severe mode collapse and instability. Therefore, DriftXpress uses a hybrid approach: projected attraction for efficiency and exact repulsion for stability.

4. Compositional Summaries (Sharding)

To handle large datasets where storing summaries for all classes in memory is infeasible, DriftXpress employs a sharding strategy. The training data is partitioned into shards (e.g., by class), and separate summaries are maintained for each shard. The final attractive field is computed by summing the contributions from each shard sequentially. This allows the method to scale to datasets like ImageNet without materializing the full summary tensor in GPU memory, trading a slight throughput reduction for memory feasibility.

Key Contributions

1. RKHS Recasting: The authors recast drifting fields within the RKHS induced by the drifting kernel, enabling low-rank, landmark-based evaluation of the attraction field.
2. DriftXpress Algorithm: They introduce a practical algorithm that replaces repeated exact attraction with cached, precomputed summaries and composes these summaries across shards for large-scale datasets.
3. Theoretical Guarantees: The paper provides bounds on the distortion of the projected field, showing that the error is controlled by the landmark kernel residual and the local kernel mass. This quantifies the conditions under which the approximation remains accurate.
4. Empirical Validation: Extensive experiments demonstrate that DriftXpress preserves sample quality while drastically reducing training time.

Results

The authors evaluate DriftXpress on SVHN, CIFAR-10, CIFAR-100, and ImageNet, comparing it against standard Drifting Models.

• Training Efficiency: DriftXpress achieves substantial speedups in wall-clock training time and throughput.
  • On CIFAR-10, it achieves a 6.63× speedup in throughput (from ~2,463 to ~16,342 images/sec) while maintaining comparable or slightly better FID scores (5.52 vs 5.64).
  • On ImageNet, the sharded variant achieves a 2.64× speedup (5,453 vs 2,058 images/sec) with comparable FID (9.21 vs 8.83).
• Convergence: DriftXpress reaches low-FID regimes faster in wall-clock time. Because the projected field uses a summary of the full dataset rather than noisy mini-batch interactions, the training dynamics are smoother, leading to faster convergence.
• Sample Quality: The method maintains sample quality comparable to standard drifting. Qualitative comparisons show that generated images are visually indistinguishable from those produced by the standard method.
• Landmark Selection: Experiments show that simple random per-class landmark selection is highly competitive and computationally efficient, often outperforming more complex selection strategies like k-center.
• Inference: Crucially, the inference cost remains identical to standard drifting (single forward pass), as the approximation is only applied during the training field estimation.

Significance and Claims

The paper claims that DriftXpress successfully pushes the training-inference trade-off of drifting models further. By moving the computational burden of field estimation from the training loop to a one-time preprocessing stage (and efficient low-rank lookups), the method reduces the primary bottleneck of drifting models—expensive training—without compromising the core advantage of one-step inference.

The authors emphasize that this approach allows drifting models to scale to larger datasets (like ImageNet) and larger batch sizes without the prohibitive memory and compute costs associated with exact kernel interactions. The work suggests that the "training cost" of one-step generative models can be significantly decoupled from the "inference cost," making high-fidelity one-step generation more practical. The paper concludes that while approximating repulsion remains an open challenge, the hybrid approach of projected attraction and exact repulsion offers a robust and scalable path forward.