LocoGen

Locality-controlled guidance for masked discrete diffusion of regulatory DNA. Without any fine-tuning, LocoGen samples enhancer sequences from a frozen MDLM (Sahoo et al., 2024) backbone using a V-REx-trained masked separator to localize the guidance gradient to invariant regulatory positions.

This repository contains the ode needed to (a) train the masked separator and (b) sample with LocoGen on top of a DRAKES MDLM backbone.

Repository layout

code/
├── pyproject.toml
├── sequence_generation/         # core package
│   ├── model/
│   │   ├── masked_separator.py  # V-REx separator + Y-head (CNN over soft one-hot)
│   │   ├── denoising_classifier.py
│   │   ├── property_regressor.py
│   │   ├── property_scorer.py
│   │   └── layers.py
│   ├── utils/
│   │   ├── train_utils.py
│   │   ├── eval_utils.py
│   │   ├── flow_utils.py
│   │   ├── grelu_compat.py
│   │   └── esm.py
│   └── datasets/
│       ├── enhancer_gosai_dataset.py
│       └── enhancer.py
├── scripts/
│   ├── masked_separator_trainer.py            # Stage 1: V-REx separator training
│   └── drakes_backbone/
│       └── mood_guided_sample_mdlm.py         # Stage 2: LocoGen sampling
└── configs/
    ├── enhancer_gosai.yaml                            # base data config
    ├── enhancer_gosai_masked_separator.yaml           # HepG2 separator
    ├── enhancer_gosai_masked_separator_k562.yaml      # K562 separator
    └── enhancer_gosai_masked_separator_sknsh.yaml     # SK-N-SH separator

Install

pip install poetry
poetry install
poetry shell

External dependencies (one-time)

LocoGen runs on the DRAKES MDLM backbone and uses DRAKES-provided reward oracles. You need to fetch them separately:

1. DRAKES code

   git clone https://github.com/ChenyuWang-Monica/DRAKES.git
   export LOCOGEN_DRAKES_REPO=$(pwd)/DRAKES/drakes_dna
   

2. DRAKES data and pretrained checkpoints (mdlm/outputs_gosai/pretrained.ckpt, reward_oracle_ft.ckpt, reward_oracle_eval.ckpt). Place them under a directory containing data_and_model/ and export:

   export LOCOGEN_DRAKES_DATA=/path/to/DRAKES_data/data_and_model/
   

3. Gosai enhancer dataset (gosai_all.csv) — drop in ./data/gosai_all.csv (the path the configs default to). See the DRAKES paper for download instructions.

The two env vars (LOCOGEN_DRAKES_REPO, LOCOGEN_DRAKES_DATA) must be set before running any of the scripts/drakes_backbone/* entry points.

Stage 1 — Train the V-REx masked separator

The separator learns a binary mask over input positions such that the masked sequence remains predictive of activity invariantly across environments (K-means clusters in the separator's x_en embedding). The HepG2 default config uses three loss terms: invariance (V-REx penalty), L_reg (sparsity), and L_smooth (mask smoothness).

python -m scripts.masked_separator_trainer \
    --config ./configs/enhancer_gosai_masked_separator.yaml \
    --out_dir ./runs/masked_separator

Replace the config to train cell-aware separators (*_k562.yaml, *_sknsh.yaml). Checkpoint of masked separator stored in ./runs/masked_separator/masked_separator_best_epoch.ckpt.

Stage 2 — Sample with LocoGen

The main entry point. λ controls the OOD push; --alpha_ratio controls the property guidance strength; the locked reference configuration is λ = 0.7, r = 3.0.

python -m scripts.drakes_backbone.mood_guided_sample_mdlm \
    --name locogen_hepg2_lam07_r3 \
    --use_masked_separator true \
    --lam 0.7 \
    --alpha_ratio 3.0 \
    --num_samples 640 \
    --batch_size 64 \
    --seed 0

The output directory contains sequences.csv (one row per generated 200 nt sequence) and a small JSON with the run config.