MUGO: Differentiable Combinatorial Optimization for Causal Variant Discovery in the Non-coding Genome
Abstract
Deciphering how non-coding variants perturb gene regulation is central to translating GWAS loci into mechanism, yet existing prioritization methods rarely deliver cell-type-resolved molecular effects, causal variant-to-gene attribution, or principled reasoning about combinatorial interactions. We introduce MUGO (Multi-head Genomic Optimization), an in silico perturbation framework that casts variant discovery as differentiable combinatorial optimization over genomic sequence. MUGO relaxes discrete edits into a continuous probabilistic nucleotide mask and performs gradient-based optimization in input space to identify single- or multi-variant perturbations that maximize a user-specified molecular objective under a sequence-to-signal foundation model. This formulation makes genome-scale search computationally tractable while retaining direct, cell-type-specific molecular readouts and enabling precise quantification of non-additive interaction effects. Across five modalities and seven tissues using two foundation-model backbones, MUGO consistently outperforms three baselines in both optimization efficiency and effect modulation, while preserving robustness and cell-type specificity. Finally, MUGO-prioritized variants are enriched for GWAS signals across diverse tissues and a broad spectrum of complex traits, turning foundation-model predictions into scalable, cell-type-resolved hypotheses for causal variant discovery and combinatorial regulatory mechanisms. Code and documentation are available at https://github.com/aicb-ZhangLabs/MUGO.