xwang112358/knitnet-rnacentral-pretrain

KnitNet — RNAcentral 37M pretrained seeds (RNA inverse design)

Masked discrete-diffusion (MDLM) generators for RNA inverse design (structure → sequence): a 37M-param GraphGPS denoiser, conditioned on a target secondary structure, recovers masked nucleotides; at inference it samples sequences that fold back to the conditioning structure. Pretrained on RNAcentral (27.5M sequences, MMseqs2 80%-identity cluster split). Intended as a pre-RL / pre-fine-tuning seed for post-training.

Objective: ELBO-weighted CE on masked tokens + a RibonanzaNet fold-back pairing loss.

Checkpoints (one relmcc-best seed per arm)

Two feature arms. Each ships the single checkpoint with the highest smoothed validation avg_relative_mcc (the fold-back selection metric — chosen on validation, not on the OOD test sets):

arm	file	dims (node/edge)	epoch	val/loss	val relmcc	features
union	union/epoch_016-periodic.ckpt	36 / 18	16	0.836	0.687	baseline + RWSE k=16 + geometric scalars + topological edges
energ	energ/epoch_032-periodic.ckpt	52 / 18	32	0.784	0.695	union + energetic loop/helix sizes

Also included: foldback/RibonanzaNet-SS.pt (fold-back model — scoring / RL reward), benchmarks/data.zip (external OOD 2D-design benchmarks), MANIFEST.json (sha256, metrics, code tag).

Which seed?

Both are viable RL seeds. On the held-out best-of-96 evaluation (below), energ leads union on graded metrics across all datasets (SR@0.8 5/5, MCC 4/5), so it's the recommended seed; union is the lighter alternative (+16 fewer node dims) and was the prior default. Most per-dataset gaps are within the target-level binomial bar (single seed), so the choice is not yet decisive.

Held-out evaluation (best-of-96 fold-back, unconstrained)

Target-level best-of-96 on 4 external OOD benchmarks + a cluster-disjoint RNAcentral test set. Values are mean ± std over each arm's top-3 checkpoints (a stability band, not a CI; single seed) — the shipped weight is the single relmcc-best of those three. SR@0.8 = fraction of targets whose best candidate reaches Jaccard ≥ 0.8; Solved = Jaccard = 1.0.

arm	dataset	Solved (J=1.0)	SR@0.8	Jaccard-best	MCC@0.8
union	openknot	0.374 ± 0.037 (~21)	0.906 ± 0.027	0.932 ± 0.005	0.930 ± 0.018
energ	openknot	0.392 ± 0.020 (~22)	0.959 ± 0.010	0.940 ± 0.003	0.918 ± 0.010
union	das	0.020 ± 0.000 (2)	0.782 ± 0.046	0.834 ± 0.005	0.463 ± 0.006
energ	das	0.020 ± 0.000 (2)	0.816 ± 0.047	0.840 ± 0.005	0.500 ± 0.018
union	eterna	0.237 ± 0.031 (~24)	0.537 ± 0.006	0.762 ± 0.003	0.687 ± 0.029
energ	eterna	0.230 ± 0.020 (~23)	0.577 ± 0.051	0.771 ± 0.008	0.727 ± 0.040
union	rnasolo	0.091 ± 0.002 (~23)	0.467 ± 0.048	0.790 ± 0.009	0.481 ± 0.038
energ	rnasolo	0.092 ± 0.004 (~23)	0.479 ± 0.050	0.801 ± 0.006	0.560 ± 0.056
union	rnacentral-1k	0.247 ± 0.008 (~247)	0.627 ± 0.005	0.824 ± 0.004	0.678 ± 0.012
energ	rnacentral-1k	0.255 ± 0.003 (~255)	0.644 ± 0.013	0.831 ± 0.002	0.696 ± 0.009

Usage (all inside the code repo)

Weights are tied to a specific code commit — tag pretrain-seed-v1 on github.com/antonia-panescu/knitnet. Work from inside the checked-out repo:

git clone git@github.com:antonia-panescu/knitnet.git knitnet-rnacentral-pretrain
cd knitnet-rnacentral-pretrain
git checkout pretrain-seed-v1            # code that loads/uses these weights

conda env create -f environment.yml && conda activate knitnet
python -m pip install -e .

# pull the weights INTO the repo (creates ./weights)
hf download xwang112358/knitnet-rnacentral-pretrain --local-dir weights
cp weights/foldback/RibonanzaNet-SS.pt external/Struct2SeQ/RibonanzaNet-SS.pt
unzip -o weights/benchmarks/data.zip -d data     # -> data/<benchmark>/…

Generate best-of-96 for a benchmark (from the repo root):

python scripts/eval/evaluate_model_MDLM.py \
    --checkpoint weights/energ/epoch_032-periodic.ckpt \
    --arm energ \                      # MUST match the checkpoint's arm (union -> --arm union)
    --data-dir data --data-subdir eterna100_v2 --dataset-name eterna100_v2 \
    --n-samples-total 96 \             # default 3-strategy preset (T0.8/16, T1.0/32, T1.0/64)
    --folding-engine ribonanza \
    --output outputs/eval/energ_ep32

Programmatic (mirror evaluate_model_MDLM.py): load_mdlm_model(ckpt) (dims auto-detected) → RNADataModule(...) with the matching arm features (union: add_rwse, add_topo_edges, add_geometric_scalars; energ: those + add_energetic_sizes) → generate_samples_batch(...).

See docs/checkpoints_for_collaborators.md in the repo for the full guide.

Gotchas

• --arm must match the checkpoint (union → --arm union, energ → --arm energ). Model dims auto-detect, but the datamodule featurization does not — a mismatch silently feeds wrong node features.
• Load .ckpt with weights_only=False (carries OmegaConf); the training arm is recorded in datamodule_hyper_parameters.
• use_cache=false for featurization (the PyG cache is not feature-aware).
• 512-nt length cap; fold-back requires RibonanzaNet-SS.pt.