chuankai-dai/utr-diffusion-checkpoint

UTR-Diffusion MCML checkpoint

Final pretrained checkpoint for UTR-Diffusion, a conditional diffusion model for multi-objective and constrained 50-nt UTR / UTR–CDS junction design.

Released file

File	Training epoch	Size (bytes)	SHA-256
checkpoints/mcml_epoch_2000.pt	2000	1,377,468,724	7125d9aa94ac67a71801c26364ce9df517a150a2e19b282caf317c835ae73ba0

This is the current MCML checkpoint trained from real_MRL_pred_MFE_260k.csv and the corresponding masked-label subsets. The pre-existing checkpoints/MRL_MFE_967k_ep_2k_ts_200_beta_0.01_cond_1_uncond_0.2_drop_0.2_lr_1e-4_at_2000epoch.pt file is a legacy CML checkpoint, is not loaded by the current examples, and is retained only for backward compatibility.

The checkpoint contains model, ema_model, optimizer, and epoch entries. The repository CLI uses the raw model state by default because the manuscript Evaluation 3 and Benchmark 2 scripts used that state. The EMA state remains selectable with --checkpoint-weights ema.

Architecture and training data

The model is the final masked continuous multi-label (MCML) configuration trained by src/scripts/train_mcml.py:

• UNet_Masked_Continuous_Multi_Labels, base and initial dimensions 200, channel multipliers (1, 2, 4), learned sinusoidal dimension 18, sequence length 50, dropout 0.2, and two per-label embeddings;
• Diffusion_Masked_Continuous_Multi_Labels, 200 timesteps, final beta 0.01, condition weight 4.0, and unconditional proportion 0.2; and
• labels real_MRL and pred_MFE.

The primary complete-label training table is real_MRL_pred_MFE_260k.csv. Training also uses its missing-MFE and missing-MRL subsets for masked multi-label learning. The datasets are not included in this model repository.

Use with the repository

git clone https://github.com/sato-lab-org/utr-diffusion.git
cd utr-diffusion
conda env create -f environment.yaml
conda activate utr-diffusion

mkdir -p checkpoints
wget -O checkpoints/mcml_epoch_2000.pt \
  https://huggingface.co/chuankai-dai/utr-diffusion-checkpoint/resolve/main/checkpoints/mcml_epoch_2000.pt

python design_utr.py \
  --mrl 8.0 \
  --mfe -2.0 \
  --out outputs/mrl_mfe_demo.fasta \
  --device cuda:0

Specify either MRL or MFE alone to use the checkpoint's masked single-label conditioning:

python design_utr.py \
  --mrl 8.0 \
  --out outputs/mrl_only_demo.fasta \
  --device cuda:0

python design_utr.py \
  --mfe -20.0 \
  --out outputs/mfe_only_demo.fasta \
  --device cuda:0

The sequence constraints --nucleotide, --amino, and --cds-amino are mutually exclusive. Without one of them, generation is conditioned only on the requested MRL and/or MFE labels.

An exact nucleotide constraint accepts arbitrary-length DNA or RNA segments:

python design_utr.py \
  --mrl 4.0 \
  --mfe -20.0 \
  --nucleotide 8:CGCTCA 32:UCA \
  --out outputs/nucleotide_demo.fasta \
  --device cuda:0

Sparse amino-acid constraints use zero-based nucleotide start positions and do not need to be contiguous or share a reading frame:

python design_utr.py \
  --mrl 8.0 \
  --mfe -2.0 \
  --amino 26:M 31:D 37:L \
  --out outputs/amino_demo.fasta \
  --device cuda:0

For a contiguous coding suffix, --cds-amino takes the complete peptide, including its initial methionine. If its length is N, the start is computed as 50 - 3N, and the peptide fills the remainder of the 50-nt sequence. The first codon is fixed to AUG. The peptide must fit (3N <= 50), and CAI control requires at least one downstream residue (N >= 2). Add --cai to request codon-adaptiveness control:

python design_utr.py \
  --mrl 8.0 \
  --mfe -2.0 \
  --cai 0.90 \
  --cds-amino MGKVKVGV \
  --out outputs/cds_cai_090.fasta \
  --device cuda:0

Here the 8-aa suffix begins at position 26. The initiating AUG is excluded from CAI, while all downstream codons contribute. --cai is the sampler's requested codon relative-adaptiveness target, not a guarantee of exact output CAI; the CLI calculates and reports the achieved sequence CAI after generation.

MRL, MFE, and CAI are independent request flags. A CAI-only command is also supported; it uses the model's no-MRL/no-MFE path and therefore still requires --cds-amino. The 8-aa form above reproduces the manuscript Evaluation 3 layout (start 26); a 10-aa peptide automatically reproduces the Benchmark 2 layout (start 20). These layouts are documented examples, not CLI presets.

The CLI refuses to replace an existing output unless --force is supplied. The advanced --targets batch form remains available for generating several joint MRL/MFE target pairs in one invocation.

Scope

This checkpoint is intended for research use with the model definitions and configuration shipped in the linked repository. Load it only from a trusted source and verify the SHA-256 above before use.

No repository-wide license has been declared for the linked source repository. Its RePaint-derived scheduler.py and utils.py retain upstream CC BY-NC-SA 4.0 notices. License metadata is therefore intentionally left unset here pending a maintainer decision; this model card does not grant additional rights.