| --- |
| library_name: transformers |
| tags: |
| - biology |
| - RNA |
| - mRNA |
| - language-model |
| - mamba |
| license: cc-by-nc-sa-4.0 |
| --- |
| |
| # Helix-mRNA |
|
|
| Minimal HuggingFace port of [Helix-mRNA](https://huggingface.co/helical-ai/Helix-mRNA) |
| -- a hybrid Mamba2 / attention language model for full-length mRNA, trained with |
| next-token prediction on single-nucleotide tokens with a codon-start marker. |
|
|
| **This repository contains code and configuration only.** The pretrained weights |
| are gated and licensed for non-commercial use, so they are not redistributed here: |
| `from_pretrained` downloads them from |
| [`helical-ai/Helix-mRNA`](https://huggingface.co/helical-ai/Helix-mRNA). Request |
| access on that page and log in (`hf auth login` or `HF_TOKEN`) before loading. |
|
|
| ## Architecture |
|
|
| | Parameter | Value | |
| |---|---| |
| | Layers | 8 (`M+M*M+M+`: 4 Mamba2, 3 MLP, 1 attention + MLP) | |
| | Attention heads | 32 query / 8 key-value (head dim 8), causal, 1 layer | |
| | Embedding dimension | 256 | |
| | FFN hidden dimension | 1024 (SwiGLU) | |
| | Mamba2 | 16 heads x 32, state size 128, expand 2, conv kernel 4, chunk size 256 | |
| | Vocabulary size | 14 | |
| | Positional encoding | None (order comes from the causal Mamba2 blocks) | |
| | Normalization | RMSNorm (pre-norm), gated RMSNorm inside Mamba2 | |
| | Architecture | Hybrid Mamba2 / Transformer decoder | |
| | Max sequence length | 12288 | |
| | Parameters | 5.2M | |
|
|
| **Vocabulary:** `A C G U N`, `E` (start of a codon), `T` (read as `U`), and the |
| special tokens `[BOS]`, `[SEP]`/`[PAD]` (shared id 1), `[CLS]`, `[MASK]`, `[UNK]`. |
| The tokenizer appends `[SEP]` to each sequence and pads on the left. |
|
|
| ## Pretraining |
|
|
| - **Objective:** Causal language modeling (next-token prediction) |
| - **Data:** mRNA sequences, with `E` inserted before each codon of the coding region |
| - **Source checkpoint:** [`helical-ai/Helix-mRNA`](https://huggingface.co/helical-ai/Helix-mRNA) (weights stored in bfloat16) |
|
|
| ## Parity Verification |
|
|
| All 9 representation levels (embedding + 8 blocks), the final norm and the LM |
| logits were verified to be bit-exact (max abs diff = 0.00) against the upstream |
| Helical blocks run in sequence (see Implementation Notes), with eager and sdpa |
| attention, on padded and unpadded batches. Verified on an NVIDIA H100 with |
| PyTorch 2.7 / CUDA 12. On 24 human mRNAs the next-token loss is 1.26 nats per |
| token (uniform over the four bases: 1.39). |
|
|
| ## Related Models |
|
|
| See the full [Helix-mRNA collection](https://huggingface.co/collections/Taykhoom/helix-mrna-6ac574d998d9cddff2ce3f9a). |
|
|
| | Model | Parameters | Notes | |
| |---|---:|---| |
| | **[Helix-mRNA](https://huggingface.co/Taykhoom/Helix-mRNA)** | 5.2M | This model | |
|
|
| ## Usage |
|
|
| ### Embedding generation |
|
|
| ```python |
| import torch |
| from transformers import AutoTokenizer, AutoModel |
| |
| tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Helix-mRNA", trust_remote_code=True) |
| model = AutoModel.from_pretrained("Taykhoom/Helix-mRNA", trust_remote_code=True) |
| model = model.to("cuda").eval() # Mamba2 blocks need a CUDA GPU |
| |
| # Mark the start of each codon in the coding region with "E" (recommended |
| # upstream); UTRs and non-coding RNA are given as plain nucleotides. |
| cds = "AUGGCCAAGUAA" |
| sequences = ["GGAUC" + "".join("E" + cds[i:i + 3] for i in range(0, len(cds), 3)), "ACGUACGU"] |
| enc = tokenizer(sequences, return_tensors="pt", padding=True).to("cuda") |
| |
| with torch.no_grad(): |
| out = model(**enc) |
| |
| token_emb = out.last_hidden_state # (batch, seq_len, 256) |
| mask = enc.attention_mask.unsqueeze(-1).to(token_emb.dtype) |
| mean_emb = (token_emb * mask).sum(1) / mask.sum(1) # masked mean pooling |
| |
| # Intermediate layers: hidden_states[0] is the embedding, [i] the output of block i |
| out_all = model(**enc, output_hidden_states=True) |
| block4_emb = out_all.hidden_states[4] |
| ``` |
|
|
| ### Next-token logits |
|
|
| ```python |
| import torch |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| |
| tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Helix-mRNA", trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained("Taykhoom/Helix-mRNA", trust_remote_code=True) |
| model = model.to("cuda").eval() |
| |
| enc = tokenizer(["ACGUACGUEAUGEGCC"], return_tensors="pt").to("cuda") |
| with torch.no_grad(): |
| out = model(**enc, labels=enc.input_ids) |
| logits = out.logits # (1, seq_len, 14); out.loss is the mean next-token loss |
| ``` |
|
|
| ### Faster attention backends |
|
|
| ```python |
| # SDPA (PyTorch 2.0+) -- recommended for production |
| model = AutoModel.from_pretrained("Taykhoom/Helix-mRNA", trust_remote_code=True, |
| attn_implementation="sdpa") |
| |
| # Flash Attention 2 (requires flash-attn) -- fastest on long sequences |
| model = AutoModel.from_pretrained("Taykhoom/Helix-mRNA", trust_remote_code=True, |
| attn_implementation="flash_attention_2", |
| dtype=torch.bfloat16) |
| ``` |
|
|
| The Mamba2 blocks require a CUDA GPU and the `mamba-ssm` kernels |
| (`pip install mamba-ssm causal-conv1d --no-build-isolation`), as upstream; |
| `causal-conv1d` is optional. |
|
|
| ### Fine-tuning |
|
|
| Standard HF conventions. For sequence-level tasks, pool over non-padding |
| positions (masked mean) before a prediction head. Gradients flow through all |
| blocks and `gradient_checkpointing_enable()` is supported. A fine-tuned model saved with `save_pretrained` includes its own |
| weights and loads from that directory without contacting the gated repo. |
|
|
| ## Implementation Notes |
|
|
| **Blocks are chained.** The upstream Helical inference code (`helical` 3.1.3, |
| `HelixmRNAPretrainedModel.forward`) never sets `hidden_states = layer_outputs[0]` |
| inside its layer loop, so every block receives the token embeddings and only the |
| last block's output reaches the final norm. This port runs the blocks in |
| sequence, as the weights were trained: next-token loss is 1.26 nats per token |
| chained versus 4.09 for the upstream forward, worse than uniform guessing (1.39). |
| Embeddings from this port therefore differ from the upstream `helical` package. |
|
|
| The Mamba2 mixer always uses the upstream inference path (separate causal |
| convolution and chunked scan, with padded positions zeroed), also in training |
| mode, where upstream switches to a fused kernel that skips the post-convolution |
| padding mask. The attention |
| block's `eager`, `sdpa` and `flash_attention_2` backends are selectable via |
| `attn_implementation`; `output_attentions=True` falls back to eager and returns |
| post-softmax probabilities. Generation caching is not implemented. |
|
|
| ## Citation |
|
|
| ```bibtex |
| @article{wood2025_helixmrna, |
| title = {Helix-{mRNA}: A Hybrid Foundation Model For Full Sequence {mRNA} Therapeutics}, |
| author = {Wood, Matthew and Klop, Mathieu and Allard, Maxime}, |
| journal = {arXiv preprint arXiv:2502.13785}, |
| year = {2025}, |
| doi = {10.48550/arXiv.2502.13785} |
| } |
| ``` |
|
|
| ## Credits |
|
|
| Original model and code by Wood et al. (Helical). Source: |
| [GitHub](https://github.com/helicalAI/helical), |
| [`helical-ai/Helix-mRNA` on the Hub](https://huggingface.co/helical-ai/Helix-mRNA). |
| Hugging Face port maintained by Taykhoom Dalal. |
|
|
| ## License |
|
|
| The weights are licensed CC BY-NC-SA 4.0 (non-commercial use only) and gated by |
| Helical, following the original repository; contact hello@helical-ai.com for |
| commercial use. This repository distributes only code and configuration. |
|
|
