OmegaGenome

01 — Overview

A generalist teacher, a deployable expert

Genomic language models keep growing — billions of parameters trained on trillions of base pairs — while the questions biologists ask of them are usually narrow: is this sequence an enhancer, does this window carry a histone mark, where are the splice sites.

OmegaGenome asks whether that narrow competence can be moved out of the large model. A pre-trained gLM is fine-tuned on one task, frozen, and then used as a teacher for a compact student trained against the task labels, the teacher's output distribution and the teacher's penultimate representation at the same time. The result keeps most of the teacher's accuracy, beats a same-architecture network trained from scratch on every task, and costs orders of magnitude less to run.

Figure 1: overview of the OmegaGenome distillation framework, with panel A fine-tuning a gLM, panel B a compact model trained from scratch, panel C the two-stage distillation, and panel D the downstream analysis over 18 prediction tasks.

Click a panel to zoom in

03 — Results

Teacher accuracy, baseline cost

Every number here is computed from the result tables released with the code. Hover a point to read it.

teacher gLM distilled 0.12 M student from-scratch baseline (0.12 M)

Accuracy — 18-task mean test MCC

Teacher	Teacher	Distilled student (0.12 M)	From-scratch baseline
Carbon-3B	0.664	0.620	0.588
Enformer	0.663	0.620	0.588
NT-2.5B	0.662	0.627	0.588
DNABERT-2	0.649	0.607	0.588
Caduceus	0.612	0.608	0.588

Cost — one RTX 3090, and 16 CPU threads

Model	Parameters	All-18 GPU time	Peak GPU memory	One task on CPU
NT-2.5B	2.54 B	7,304 s	11,553 MB	8.3 h
Carbon-3B	3.45 B	1,073 s	8,004 MB	6.6 h
DNABERT-2	117 M	531 s	3,074 MB	23.6 min
Caduceus	7.7 M	342 s	489 MB	4.2 h
Enformer	251 M	102 s	1,184 MB	1.9 min
OmegaGenome student	0.12 M	11.1 s	17 MB	2.65 s

Beyond classification — base-resolution regression

The same recipe transfers to a 34-track, per-base-pair regression benchmark, where each model predicts a coverage value at every base of a 32,768 bp window. Distilling NTv3-650M into students of 0.12 M, 8 M and 100 M parameters traces a smooth cost–accuracy curve, and the distilled student beats a size-matched from-scratch baseline in every assay family and at every size.

Model	Parameters	Whole-test GPU time	Peak GPU memory
NTv3-650M teacher	655 M	2,122 s	9,472 MB
NTv3-100M student	106 M	705 s	3,909 MB
NTv3-8M student	7.7 M	136 s	1,210 MB
BPNet student	0.12 M	10.6 s	89 MB

04 — Figures

Main-text figures

Vector, as submitted. Click any figure to open it full size.

Nine-task lollipop grid, 18-task mean MCC and relative improvement over the baseline for five teachers and their distilled students. — **Figure 2.** Distilled experts match teacher performance across diverse benchmarks. (A) MCC on nine representative tasks for the five teachers (open rings), their distilled students (dots) and the baseline (dotted line). (B) Mean MCC over all 18 tasks. (C) Mean relative improvement of each student over the baseline.

Accuracy versus model size, GPU time and peak memory, for classification and for base-resolution regression. — **Figure 3.** Distilled students match teacher accuracy at a fraction of the compute. Top row, 18 classification tasks: mean MCC versus model size (A), total GPU time (B) and peak GPU memory (C), each bar labelled with its fold-change over the student. Bottom row, 34-track regression: mean Pearson correlation versus model size (D), whole-test GPU time (E) and peak memory (F).

Distillation-method comparison, student-size scaling, loss-term ablation and loss-weight landscape. — **Figure 5.** Robustness to the distillation method, the student size and the loss weights. (A) OmegaGenome versus DKD, DIST and logit standardisation. (B) Student-size scaling. (C) Loss-term ablation and (D) the distillation-versus-feature weight landscape, aggregated over five teachers × 18 tasks with MCC z-scored within each task.

All six main-text figures, in vector PDF and PNG, are in the repository under paper_figures/figures/, next to the scripts that draw them.

05 — Reproduce

Everything is in one repository

Figures, on a laptop

22 of the 23 figure scripts redraw the paper's panels from the bundled result tables — no GPU, no downloads, seconds per figure.

Training, on one GPU

Stage 1 fine-tunes a teacher per task; stage 2 distils the student. Both are single commands with every hyperparameter overridable.

Nothing hidden

Machine paths come from documented environment variables, dependencies are locked, and the lint workflow runs on every push.

Quick start

git clone https://github.com/aicb-ZhangLabs/OmegaGenome.git
cd OmegaGenome && cp env_sample.toml env.toml

# redraw the paper's figures from the bundled tables
uv sync --extra analysis
cd paper_figures
python figure_panels/build/fig4_efficiency_carbon.py     # Figure 3

# or train: fine-tune a teacher, then distil the student
uv sync --extra default
python -m src.train.finetune_teacher carbon_3b --config.task-names H3K4me3
python -m src.train.distill carbon-raw --task-names H3K4me3

The repository README documents every entry point, the environment variables and the released checkpoints.

06 — Citation

Cite this work

If this work is useful to you, please cite the paper and, if you use the code directly, the archived snapshot.

@software{omegagenome_code,
  title  = {OmegaGenome: source code},
  author = {Xu, Pengcheng and Liu, Junhao and Dai, Yi and Nishida, Kainoa Andrew and
            Sun, Dongbo and Lei, Yutong and Hu, Yaqi and Wang, Chaoyang and Zhang, Jing},
  year   = {2026},
  doi    = {10.5281/zenodo.22805335},
  url    = {https://doi.org/10.5281/zenodo.22805335}
}