Helix-mRNA

Minimal HuggingFace port of 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. 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 (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.

Model	Parameters	Notes
Helix-mRNA	5.2M	This model

Usage

Embedding generation

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

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

# 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

@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, helical-ai/Helix-mRNA on the Hub. 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.