Distributional Alignment and Diversity Control for Generative Models via Optimal Transport and Uncertainty Quantification
Problem and high-level solutions. Foundation models lose reliability under domain shift and often collapse diversity after alignment. We address this by steering a model’s output distribution at inference toward a target distribution while preserving task utility, and by decoding policies that maintain solution breadth without sacrificing correctness. The core mechanism is optimal transport applied on model outputs or latent features to minimize distribution mismatch. A calibration layer provides coverage guarantees and safe refusal.
Approach. We design OTbased steering maps that act on the base model output distribution so that this transported distribution approaches the aligned distribution . We combine this with diversitypreserving decoding using sequential Monte Carlo or tree search with repulsive interactions in representation space. We evaluate Wasserstein distance reduction, utility retention, and diversity metrics; we add conformal risk control for selective abstention.
Relevance and impact. Healthcare and education are immediate beneficiaries in the VinUni and Vingroup ecosystem. In healthcare, steering supports more consistent medical reporting across sites while preserving diagnostic fidelity. In education, decoding strategies increase breadth of hints and worked solutions while refusing low confidence answers. Outputs might include new algorithms with theoretical guarantee, an open library, an Aligned Diversity benchmark, and four or more toptier or Q1 publications.