LocoGen

Separator weights for LocoGen: locality-controlled guidance for masked discrete diffusion of regulatory DNA. LocoGen samples enhancer sequences from a frozen MDLM backbone, with no fine-tuning, by localizing the guidance gradient to invariant regulatory positions using a V-REx-trained masked separator.

Code: https://github.com/Lotfollahi-lab/LocoGen

What this repository contains

Only the masked separator — the component trained in LocoGen's Stage 1. One checkpoint per cell line.

File	Cell line	Params
hepg2/masked_separator.ckpt	HepG2	11,135,495
k562/masked_separator.ckpt	K562	11,135,495
sknsh/masked_separator.ckpt	SK-N-SH	11,135,495

Each is a plain {"model": state_dict} payload (316 tensors, ~45 MB).

What it does not contain

LocoGen runs on the DRAKES MDLM backbone and uses the DRAKES reward oracles. Those are a separate release and are not redistributed here. You need them to sample anything. Fetch both as described in the LocoGen README:

• mdlm/outputs_gosai/pretrained.ckpt — frozen MDLM backbone
• mdlm/outputs_gosai/lightning_logs/reward_oracle_ft.ckpt — guidance oracle
• mdlm/outputs_gosai/lightning_logs/reward_oracle_eval.ckpt — held-out scoring
• mdlm/gosai_data/processed_data/gosai_all.csv — Gosai dataset

Architecture comes from the config, not the checkpoint

The loader builds the model from a YAML config and then loads these weights into it. The checkpoint carries no architecture metadata, so you must pass the matching config from the code repository:

Checkpoint	--ms_config	target_idx
hepg2/masked_separator.ckpt	configs/enhancer_gosai_masked_separator_hepg2.yaml	0
k562/masked_separator.ckpt	configs/enhancer_gosai_masked_separator_k562.yaml	1
sknsh/masked_separator.ckpt	configs/enhancer_gosai_masked_separator_sknsh.yaml	2

At inference only the architecture fields and target_idx are read; the training-side loss weights in those configs do not affect sampling.

Usage

Download a checkpoint and point the LocoGen sampler at it:

export LOCOGEN_DRAKES_REPO=/path/to/DRAKES/drakes_dna
export LOCOGEN_DRAKES_DATA=/path/to/data_and_model/

CKPT=$(python -c "from huggingface_hub import hf_hub_download; \
print(hf_hub_download('DabinJeong/LocoGen', 'hepg2/masked_separator.ckpt'))")

Then run Stage 2 sampling. These are the full reference settings — --lam and --alpha_ratio alone are not sufficient to reproduce the reported numbers:

python -u -m scripts.drakes_backbone.mood_guided_sample_mdlm \
    --name locogen_hepg2 \
    --cell_type hepg2 \
    --ms_config ./configs/enhancer_gosai_masked_separator.yaml \
    --ms_ckpt "$CKPT" \
    --lam 0.7 \
    --alpha_ratio 3.0 \
    --alpha_cap 20.0 \
    --ood_sil_localized true \
    --ood_sil_threshold 0.5 \
    --sil_t_threshold 0.3 \
    --num_samples 640 \
    --batch_size 64 \
    --seed 0

--num_steps is left at its default of 128, as in the original runs.

For K562 or SK-N-SH, swap --cell_type, --ms_config and --ms_ckpt together — the oracle head is selected by --cell_type, so a mismatched pair silently scores against the wrong cell line.

The run writes sequences.csv (one row per generated 200 nt sequence) and a summary.json. Reported results average over seeds 0, 1 and 2.

Loading the weights directly, without the sampler:

import torch
from huggingface_hub import hf_hub_download
from omegaconf import OmegaConf
from locogen.model.masked_separator import MaskedSeparatorModel

ckpt = hf_hub_download("DabinJeong/LocoGen", "hepg2/masked_separator.ckpt")
cfg = OmegaConf.load("configs/enhancer_gosai_masked_separator.yaml")
model = MaskedSeparatorModel(cfg)
model.load_state_dict(torch.load(ckpt, map_location="cpu")["model"])
model.eval()

# x: (B, 200, 4) soft one-hot -> M: (B, 200) mask in [0, 1]
M = model.separator(x)

Citation

LocoGen, Dabin Jeong. See the code repository for the current reference and publication status.