Locality-Controlled Out-of-Distribution Guidance for Robust Regulatory DNA Sequence Design

AI summary

Overview

Research area: Machine learning for biological sequence design — specifically, property-conditioned generative modeling of regulatory DNA (enhancers) using discrete diffusion models with inference-time guidance.

Technical level: Advanced. The paper assumes familiarity with diffusion models, masked discrete diffusion language models (MDLM), classifier guidance, score-based OOD control, and invariant risk learning.

Scope (one sentence): The paper diagnoses how existing guidance methods for enhancer design "reward-hack" the activity predictor, and proposes LocoGen, a frozen-backbone sampling-time method that restricts out-of-distribution pressure to background sequence positions so that property, novelty, and plausibility are satisfied simultaneously.

What This Paper Is About

Generative models can propose regulatory DNA sequences predicted to drive high gene expression, but scoring them only on predicted activity hides failure modes: the model may simply stuff repeated motifs into a sequence or drift toward a generic GC-rich composition, both of which fool the predictor without producing realistic enhancers. The authors argue that a good candidate must simultaneously achieve high predicted activity, be novel relative to the training set, and preserve the natural "grammar" of regulatory DNA (motif diversity and density). Their goal is to steer a pretrained diffusion model into the region of design space where all three hold at once — a region they call the "empty corner," which no existing baseline reaches.

Key Contributions

1. A three-axis diagnostic framework. The authors decompose evaluation of generated regulatory DNA into property (predicted activity and ATAC accessibility), novelty (authenticity and rarity), and compositional plausibility (deviation in motif-usage Gini index and motif hits per sequence, written ΔGini and Δhits). This decomposition exposes two distinct reward-hacking modes that single-axis evaluation conceals.

2. Diagnosis of two structurally different hacking modes. SVDD inflates motif density to 1.44× the natural baseline (Δhits = 0.248, from 0.561 to 0.809) while concentrating motif usage (ΔGini = 0.231, from 0.308 to 0.539); DRAKES instead drifts toward a GC-rich background with skewed motif usage (ΔGini = 0.261) but little density inflation (Δhits = 0.06).

3. Locality-controlled guidance (LocoGen). Guidance is reformulated as a question of where to perturb rather than how much. A learned soft mask from a "masked separator" identifies invariant (motif-like) versus environment (background) positions, and OOD damping is applied only to environment positions during sampling, on a frozen MDLM backbone with no retraining.

4. An invariant-learning mask module. The separator is trained with a sufficiency loss, a variance risk extrapolation (V-REx) invariance penalty across E = 3 k-means-discovered environments, a mask-density regularizer, and a fused-lasso total-variation smoothing penalty, producing sparse, contiguous motif masks.

Main Findings

• LocoGen fills the empty corner in HepG2. Pred-activity 4.679 ± 0.021, ATAC_acc 0.716 ± 0.019, authenticity 0.906 ± 0.006 (the best of any method), rarity +0.71 ± 0.00, |ΔGini| 0.055 ± 0.005, |Δhits| 0.057 ± 0.003. The authors note it is not best on every metric (SVDD reaches 6.830 predicted activity, DRAKES +1.08 rarity) but is the property-optimizing method that stays acceptable on all three axes.

• Single-axis optimization produces clear trade-offs. SVDD (pred-activity 6.830 ± 0.024) and DRAKES (5.610 ± 0.070) lead on activity but show large plausibility shifts; DNA-Diffusion nearly matches training composition (|ΔGini| 0.002 ± 0.002, |Δhits| 0.010 ± 0.002) but has weaker property optimization (2.595 ± 0.087) and negative rarity (−0.07 ± 0.03).

• Classifier guidance is the main driver of property optimization. Removing guidance (γ = 0 rows) drops predicted activity to roughly 0.157–0.309; enabling classifier guidance alone lifts it to 4.516 ± 0.151.

• Localization gives a modest but consistent gain. Replacing learned local OOD with global OOD at random positions yields pred-activity 4.608 ± 0.061 versus 4.679 ± 0.021, with novelty and plausibility also slightly worse.

• Adversarial localization hurts. Deliberately reversing property guidance at separator-identified invariant positions ("Attack") reduces pred-activity to 3.881 ± 0.038, supporting the claim that the mask finds activity-sensitive positions.

• Localization alone cannot replace generative optimization. In a non-generative segment-and-shuffle control, the learned mask beat a GC-content heuristic (predicted activity 0.442 vs 0.359), but both were far below full LocoGen (4.679; ATAC accessibility 0.280 vs 0.716).

• The separator learns a meaningful partition. The invariant subsequence stays close to the natural baseline, is enriched for activity-associated JASPAR motifs, and has larger functional impact under ablation than random regions.

• Results hold under alternative evaluation configurations. Two alternative publicly released oracles (from Li et al., 2024 and Genentech, 2025) largely preserve the ranking, and a hyperparameter sweep over each method's published ranges retains LocoGen's balanced profile.

Methodology in Plain English

The authors start from an existing pretrained masked discrete diffusion model, which generates DNA by starting from a fully masked sequence and progressively filling in nucleotides. At sampling time, three forces act on each position: the model's own prior over nucleotides, a pull toward higher predicted activity (classifier guidance), and an out-of-distribution (OOD) damping term that flattens the distribution toward uniform to encourage exploration beyond the training set.

Their key change is to make the OOD damping position-specific. A separate network — the masked separator — is trained to output a soft per-position mask that splits a sequence into an "invariant" part that reliably predicts activity across subpopulations, and an "environment" part that does not. To get this, they train a predictor on the invariant part only, and penalize variance in prediction error across environments discovered by clustering the environment encoder's representations (with the environment encoder itself frozen and the clusters refit each epoch). Sparsity and smoothness penalties push the mask toward contiguous blocks. The mask is then applied during sampling: positions marked as background get the OOD damping factor 1 − √γ, but only in the latter half of the sampling trajectory, when the model's estimate of the clean sequence is reliable. Motif positions are left untouched by damping.

The strength of the activity pull is set adaptively at each step so that its magnitude stays a fixed fraction of the prior's magnitude, and the gradient is added only to the four nucleotide channels while preserving the model's prescribed unmasking schedule. The backbone is never retrained.

Evaluation uses approximately 700k 200-bp human enhancer sequences with expression measurements in HepG2, K562, and SK-N-SH, following the DRAKES protocol and chromosomal split. Each method generates N = 640 sequences (20 batches × 32), which are rescored by an independent held-out oracle with three MPRA readouts plus a binary ATAC classifier. Metrics are reported as mean ± standard deviation over three sampling seeds. Baselines are MDLM, MDLM+CG, DNA-Diffusion, DRAKES, and SVDD, each run with authors' released checkpoints and original decoding hyperparameters.

Why This Matters

Research impact. The paper reframes guidance strength as a spatial rather than scalar control problem, which is a transferable idea: any sequence modality with non-uniform positional constraint (promoters, RNA, protein) could in principle be guided this way. The three-axis evaluation also provides a template for auditing reward hacking in biological generative models, where a high predictor score is otherwise easy to mistake for success.

Real-world applications:

• Cell-type-specific enhancer design for gene therapy vectors, where candidate sequences must be both potent and compositionally natural to function in cells.
• Synthetic biology and biomanufacturing, where tuning expression of a target gene requires regulatory elements that work reliably outside the predictor.
• Reporter assay and MPRA experimental design, where pre-filtering candidates by plausibility saves wet-lab budget.
• Safety and screening pipelines, where detecting compositionally aberrant designs matters as much as ranking them by predicted activity.

Industry relevance. Because LocoGen runs on a frozen backbone with no retraining, a new target property or cell line requires only a swap of the reward model and separator, not a new fine-tuning campaign — an important practical advantage over reward-based fine-tuning approaches such as DRAKES, which require re-specifying the training objective and a separate fine-tuning pass per objective.

Future Directions

• Wet-lab validation. All reported evidence is computational (held-out oracles and an ATAC classifier). The introduction frames design as lab-in-the-loop, but no experimental validation is reported; testing LocoGen candidates in an MPRA would be the decisive next step.
• Head-to-head comparison with concurrent methods. GLID2E, Ctrl-DNA, and GFlowNet are explicitly excluded because no public checkpoints are available and retraining was out of scope, leaving their comparative positioning open.
• Extension beyond activity and beyond HepG2/K562/SK-N-SH. The separator is trained for a specific activity signal; whether the same invariant mask transfers across properties, cell lines, or sequence lengths is not established.
• Porting locality control to other backbones. The method is instantiated on masked discrete diffusion (MDLM). Whether localized OOD damping transfers to continuous diffusion or flow-matching backbones, or to autoregressive models, is untested.

Target Audience

Machine learning researchers working on diffusion models and inference-time guidance, particularly those interested in reward hacking and OOD control; computational biologists and genomics researchers designing regulatory elements; and practitioners in gene therapy, synthetic biology, and biomanufacturing who need generative sequence design that holds up beyond a single predictor score. Readers without a background in diffusion models or sequence biology will find the diagnostic framing accessible, but the methodology sections require intermediate-to-advanced familiarity.

Authors’ abstract

Designing regulatory DNA sequences for targeted gene-expression control is an important challenge. Recent approaches frame this task as property-conditioned generation, where diffusion models trained on natural sequences are guided toward high predicted activity through reward-based fine-tuning or sampling-time guidance. We show that single-reward evaluation can obscure reward-hacking behavior. By decomposing evaluation into property, novelty, and plausibility, we identify an empirically underexplored region of the evaluation space in which all three objectives are satisfied, but no existing baseline succeeds. We postulate that this gap arises because prior methods treat guidance primarily as a matter of magnitude. In regulatory DNA, however, sequence grammar is spatially nonuniform: motif positions are sensitive to perturbation, while background regions allow more exploration. We therefore propose LOCOGEN by reformulating guidance as a problem of locality. LOCOGEN learns an invariant motif mask through sequence-level invariance learning to protect motif positions, while applying OOD damping only to background positions. It operates at sampling time on a frozen diffusion backbone, without retraining the backbone. Across HepG2, K562, and SK-N-SH, LOCOGEN reaches this previously underexplored region of the property-novelty-plausibility space. Code is available at https://github.com/Lotfollahi-lab/LocoGen.