MIMIC / README.md
jwkova's picture
Upload README.md with huggingface_hub
27688f6 verified
|
Raw
History Blame Contribute Delete
9.29 kB
---
license: mit
thumbnail: https://raw.githubusercontent.com/PolymathicAI/MIMIC/main/assets/MIMIC_logo.png
tags:
- biology
- genomics
- dna
- rna
- protein
- multimodal
- foundation-model
datasets:
- polymathic-ai/LORE-examples
---
<!--
TODO before making the repo public:
- Fill in the Training data section.
-->
<p align="center">
<img src="MIMIC_logo.png" alt="MIMIC" width="320">
</p>
# MIMIC 1.0
MIMIC is a multimodal encoder–decoder foundation model of the central dogma, trained
jointly over **DNA, RNA, and protein** together with a range of structural and
functional tracks. A single model embeds any subset of modalities into a shared
representation space and generates any modality conditioned on the others.
- **Architecture:** 20-layer encoder × 1536-d, 12-layer decoder × 1536-d, rotary
position embeddings, mixed (uni-/bi-directional) attention, 5 register tokens.
- **Parameters:** ~1.25B.
- **Weights:** `bfloat16`, `safetensors`.
- **Modalities (26):** nucleotide/amino-acid sequence, codons, splice junctions/regions,
coding annotation, conservation (phyloP human/mouse), protein structure tokens, DSSP,
SASA, MaSIF surface features, protein abundance, RASP2 reactivity, and free-text
functional/context channels.
## Benchmarks & Transfer Learning
MIMIC's learned representations transfer to RNA and protein property prediction, where a single multimodal model is competitive with or ahead of strong single-modality foundation models:
- **Protein — [PFMBench](https://arxiv.org/abs/2604.24506)** (function, structure, interaction, and developability): MIMIC matches or exceeds protein-only baselines including **ESM3, ESM-C, ProTrek, and SaProt**, leading on at least **7 of 11** tasks — with especially strong protein–ligand binding (BindingDB, PDBbind) and the best results across all developability tasks.
- **mRNA — [mRNABench](https://github.com/morrislab/mRNABench)** (function, localization, transcriptional regulation, and variant-effect prediction): MIMIC outperforms **Evo 2** and **Orthrus** on **4 of 7** tasks and a dilated-ResNet baseline on **6 of 7**, leading overall.
Beyond property prediction, the same model supports conditioned design — e.g. identifying corrective edits for a clinically relevant HBB splice-disrupting mutation, and generating diverse, high-confidence binder sequences by jointly conditioning on backbone shape and surface chemistry (PD-L1, hACE2).
See the [paper](https://arxiv.org/abs/2604.24506) for the full benchmark tables and comparisons.
## Usage
```python
from mimic import load_pretrained
# Downloads config.json + model.safetensors pinned to git tag v1.0.
model = load_pretrained(version="1.0")
# Real matched examples from LORE (raw view = modality name -> raw value).
from datasets import load_dataset
ex = load_dataset("polymathic-ai/LORE-examples", split="train") # default config = raw
prot = next(r for r in ex if r["kind"] == "protein") # has aa_seq + sasa
rna = next(r for r in ex if r["kind"] == "rna") # has rna_seq + splice_jctns_5cls
# --- Embed ---
model.input([{"rna_seq": rna["rna_seq"]}])
reps = model.embed() # {"full": [B, N, 1536], "mod_ids": [B, N] per-token group ids}
# reps["full"] -> (batch, num_tokens, 1536); pass return_register=True / return_modality=True for more
# --- Generate demo 1 (cross-modal): protein sequence -> per-residue solvent accessibility ---
# The target length is inferred from the co-grouped aa_seq; the default strategy is an
# Ensemble soft-vote (deterministic at low temp).
model.input([{"aa_seq": prot["aa_seq"]}])
out = model.generate("sasa")
print(out["sasa"]) # per-residue SASA (float array)
# --- Generate demo 2 (cross-modal): RNA sequence -> per-position splice-site classes ---
model.input([{"rna_seq": rna["rna_seq"]}])
out = model.generate("splice_jctns_5cls")
print(out["splice_jctns_5cls"]) # one 5-class site label per position
# --- Richer output: raw arrays alongside the detokenized prediction ---
# By default generate() returns just the detokenized prediction. Set any of
# return_tokens / return_logits / return_probs / return_sampling_probs and each
# value becomes a dict with "preds" plus the extras you asked for (numpy arrays).
model.input([{"rna_seq": rna["rna_seq"]}])
out = model.generate("splice_jctns_5cls", return_probs=True, return_tokens=True)
print(out["splice_jctns_5cls"]["preds"]) # per-position class labels (as above)
print(out["splice_jctns_5cls"]["probs"].shape) # (num_positions, num_classes)
print(out["splice_jctns_5cls"]["tokens"].shape) # (num_positions,) predicted token ids
```
`load_pretrained` fetches only `config.json` + `model.safetensors`; tokenizers ship
inside the `mimic` package.
**Install:** `pip install git+https://github.com/PolymathicAI/MIMIC.git` (imports as `mimic`).
Source and docs: [github.com/PolymathicAI/MIMIC](https://github.com/PolymathicAI/MIMIC).
## Files
| File | Description |
|------|-------------|
| `config.json` | Architecture + modality configuration consumed by `load_pretrained`. |
| `model.safetensors` | Model weights (bf16). |
## Example data
Ready-to-run examples (linked in the sidebar under **Datasets**):
- [`polymathic-ai/LORE-examples`](https://huggingface.co/datasets/polymathic-ai/LORE-examples)
— a small set of matched multimodal rows (25 modalities) in both raw and tokenized
form, for quickstart / embedding / generation.
```python
from datasets import load_dataset
ds = load_dataset("polymathic-ai/LORE-examples")
```
## Modalities
MIMIC represents each molecule as a set of co-observed modalities grouped into
three tracks: **nucleic** (RNA/DNA and its per-position annotations), **protein**
(amino-acid sequence, structure, and derived features), and **text** (free-text /
categorical context). Pass any modality to `model.input()` under its short name
(e.g. `rna_seq`) or pre-tokenized under its `tok_` key. The authoritative,
per-checkpoint list is `model.modality_info`; the `Track` column drives pathway
gating (generation is limited to within-track and text-association pathways). A
few assay tracks (`atac`, `cage`, `rasp2`, `prot_abund`) are **context-conditional**:
pass a free-text `context` alongside them to condition on cell-state / assay
metadata — the `Conditioning context` column shows a real example for each.
See the full modality table (with per-modality examples and conditioning contexts) on the [LORE-examples dataset card](https://huggingface.co/datasets/polymathic-ai/LORE-examples).
<!-- Training data section — drafted, held back for now. Restore (uncomment) when ready.
## Training data
MIMIC 1.0 was trained on **LORE**, a corpus that aligns heterogeneous molecular data
into coherent, *partially observed* multimodal examples. Each example is organized
around a shared transcript / protein anchor, so DNA, RNA, protein, and their
associated structural, regulatory, and functional tracks are co-registered on the
same underlying biological entity. Most examples observe only a subset of the
modalities, and the model is trained to map between any observed subset and the rest
(multi-pathway training), with a context-length curriculum scaling from 1k to 10k tokens.
**Scale**
- ~13M RNA transcripts
- ~15.5M proteins
- 4B+ natural-language tokens (functional / contextual annotations)
- 6,000+ organisms
**Modalities (26).** Nucleotide and amino-acid sequence; RNA codons; splice and CDS
junctions and regions; coding annotation and feature type; sequence conservation
(phyloP, human and mouse); protein structure tokens, DSSP secondary structure,
solvent-accessible surface area (SASA), and MaSIF surface features; protein
abundance; cell/assay experimental tracks (CAGE, ATAC, and RASP2 reactivity); and
free-text functional and contextual channels.
A small, ready-to-run sample of LORE is released as
[`polymathic-ai/LORE-examples`](https://huggingface.co/datasets/polymathic-ai/LORE-examples).
See the [paper](https://arxiv.org/abs/2604.24506) for full dataset construction,
per-modality sources, and preprocessing details.
-->
## Citation
If you use this work, please cite:
```bibtex
@misc{golkar2026mimicgenerativemultimodalfoundation,
title={MIMIC: A Generative Multimodal Foundation Model for Biomolecules},
author={Siavash Golkar and Jake Kovalic and Irina Espejo Morales and Samuel Sledzieski and Minhuan Li and Ksenia Sokolova and Geraud Krawezik and Alberto Bietti and Claudia Skok Gibbs and Roman Klypa and Shengwei Xiong and Francois Lanusse and Liam Parker and Kyunghyun Cho and Miles Cranmer and Tom Hehir and Michael McCabe and Lucas Meyer and Rudy Morel and Payel Mukhopadhyay and Mariel Pettee and Helen Qu and Jeff Shen and David Fouhey and Hadi Sotoudeh and Vikram Mulligan and Pilar Cossio and Sonya M. Hanson and Alisha N. Jones and Olga G. Troyanskaya and Shirley Ho},
year={2026},
eprint={2604.24506},
archivePrefix={arXiv},
primaryClass={cs.AI},
url={https://arxiv.org/abs/2604.24506},
}
```
## License
This model and its accompanying source code are released under the [MIT License](https://github.com/PolymathicAI/MIMIC/blob/main/LICENSE).