Generative Retrieval via Diffusion Transformer with Metric-Ordered Sequence Training and Hybrid-Policy Preference Optimization

Problem Definition: Pattern-Preserving Attribute Retrieval

The paper addresses a specific limitation in production retrieval systems: the tension between satisfying a target attribute (e.g., quality, safety, policy compliance) and preserving a fine-grained pattern (e.g., a specific visual style, topic, or failure mode) defined by a seed set.

Standard approaches fail to balance these goals:

• Averaging seed embeddings preserves the pattern but often lands in a low-attribute-density region, retrieving ordinary items that lack the desired attribute.
• Global attribute retrieval finds items with the target attribute but often drifts into unrelated patterns, flooding the results with irrelevant variants.

The authors formalize this as Pattern-Preserving Attribute Retrieval. The objective is to generate a query embedding that retrieves top-$K$ neighbors which are simultaneously attribute-positive and of the seed pattern. They define a primary intersection metric, Joint@K, which penalizes both attribute failure and pattern drift.

Methodology: MO-DiT+HPPO

The proposed framework, MO-DiT+HPPO (Metric-Ordered Diffusion Transformer with Hybrid-Policy Preference Optimization), is a staged training pipeline designed to navigate the embedding space from low-density to high-density regions while staying within a specific pattern cluster.

1. Architecture: Continuous Generative Retriever

The core model is a Diffusion Transformer (DiT) operating in a frozen multimodal embedding space.

• Input: A sequence of item embeddings (the seed set) acts as condition tokens.
• Output: A continuous query embedding generated via flow matching.
• Mechanism: The model learns a velocity field to transport Gaussian noise to a target embedding along a geodesic path on the unit hypersphere. Unlike discrete generative retrieval, it synthesizes a continuous vector directly, avoiding hallucination of invalid identifiers.

2. Metric-Ordered Sequence Construction

To teach the model the direction of "improvement" without dense online labeling, the authors construct training trajectories:

• Clustering: Attribute-positive items are clustered to define latent pattern proxies.
• Sparse Labeling: A small subset of items is labeled online for "recall density" (fraction of top-$K$ neighbors satisfying the attribute).
• Prediction & Ordering: A lightweight metric predictor is trained on these sparse labels to score the entire pool. Items within a pattern cluster are sorted from low to high predicted attribute density.
• Result: Training sequences represent a trajectory moving from ordinary in-pattern items toward high-attribute-density regions.

3. Staged Training Pipeline

The framework proceeds through four stages, all using the same DiT backbone and flow-matching objective but with different data and supervision targets:

• Stage 1: Large-Scale Raw-Sequence Pretraining: Establishes a general continuous retrieval prior using natural sequences (e.g., user behavior streams) without metric ordering.
• Stage 2: Shared Multi-Domain Metric-Ordered Continuation Pretraining (CPT): Trains on the constructed metric-ordered sequences across all domains. The model learns to predict the next item in a low-to-high density trajectory, internalizing the "metric-improvement direction" as a general capability.
• Stage 3: Domain-Specific Tail-Centroid Supervised Fine-Tuning (SFT): Maps the input prefix (first $m$ items) to the centroid of the high-density tail (last $L-m$ items). Using a centroid rather than a single top item reduces sensitivity to noise and outliers.
• Stage 4: Hybrid-Policy Preference Optimization (HPPO): Aligns the model with the true online objective (Joint@K).
  • Hybrid Candidate Pool: Generates candidates from two sources: (1) deterministic tail-centroid constructions (static, pattern-preserving) and (2) policy-generated queries sampled with a "guidance-scale fan" (varying the trade-off between pattern and attribute).
  • Labeling: Candidates are scored using real online retrieval (Joint@K and Same@K).
  • Pareto Pair Filter: A critical component that filters preference pairs. A "winner" (higher Joint@K) is only accepted if it does not reduce the Same-Pattern Share (purity) relative to the "loser." This prevents the model from gaining attribute scores by drifting to unrelated patterns.
  • Optimization: Uses a reference-anchored preference loss (similar to DPO) where the score is the negative flow-matching loss, regularized by a frozen SFT reference model.

Key Contributions

1. Formalization: Defined Pattern-Preserving Attribute Retrieval and the Joint@K intersection metric.
2. Metric-Ordered Sequences: Introduced a method to convert sparse online labels into large-scale, in-pattern training trajectories that encode the direction of attribute improvement.
3. Staged Framework: Demonstrated that shared multi-domain CPT followed by domain-specific tail-centroid SFT improves retrieval across diverse domains.
4. HPPO: Adapted preference optimization for continuous diffusion retrieval, utilizing a hybrid candidate pool and a tradeoff-aware Pareto pair filter to ensure attribute gains do not come at the cost of pattern purity.
5. Evaluation Protocol: Established a strict evaluation protocol with item- and pattern-holdout splits and paired-bootstrap significance testing.

Results

The paper evaluates the approach on four large-scale attribute domains (D1–D4) under strict holdout protocols.

• Performance: The MO-DiT+HPPO pipeline significantly outperforms a strong pretrained generative retriever and simple averaging baselines.
• Significance: Gains in Joint@K are large and statistically significant (paired-bootstrap) in seven of eight domain-split cells. On the hardest domain's pattern split, the method achieves a marginal tie, indicating robustness in difficult scenarios.
• Ablations:
  • Pareto Filter: Removing the filter causes the model to slide along the attribute-pattern frontier (gaining attributes but losing purity). The filter pushes the frontier outward, improving Joint@K while maintaining or increasing Same-Pattern Share.
  • Candidate Policy: Static (tail-centroid) and policy-generated (guidance-fan) candidates are shown to be complementary; the hybrid pool yields better results than either alone.
  • Ordering: Metric-ordered training outperforms non-ordered or single-domain baselines.

Significance and Claims

The paper claims that MO-DiT+HPPO successfully solves the tension between attribute satisfaction and pattern preservation in continuous generative retrieval. By treating retrieval as a generation problem in a continuous space and using a staged approach that moves from general priors to specific metric-ordered trajectories, the system learns to synthesize queries that target high-density regions of specific patterns.

The authors emphasize that the Pareto pair filter is the key ingredient that prevents "reward hacking" (sacrificing pattern for attribute). The framework allows for practical, per-domain refreshes because the expensive pretraining is shared, while the final alignment (HPPO) is computationally light and relies on real online metrics rather than learned proxy rewards. The work bridges the gap between offline generative modeling and online retrieval objectives, providing a method to expand seed sets for downstream tasks like content discovery, variant detection, and dataset curation without losing the specific intent of the seed pattern.