Attention meets DNA: Enformer, now in PyTorch

What it does

This repository re-implements DeepMind’s Enformer, a transformer-style network that predicts gene expression from long DNA sequences, entirely in PyTorch. It bundles the official TensorFlow Sonnet weights—ported to PyTorch and hosted on HuggingFace—along with task-specific fine-tuning wrappers and a GenomicIntervalDataset helper that reads .bed files, handles chromosome name remapping, and applies shift or reverse-complement augmentations on the fly.

The interesting bit

The heavy lifting here is numerical archaeology: the author tracked down a xlogy gamma discrepancy between TensorFlow and PyTorch that was causing layer-by-layer drift, then locked in precomputed gamma positions so the ported model actually hits the paper’s reported Pearson R values. Beyond the port, the repository exposes intermediate embeddings and offers three distinct adapter wrappers—head, static context, and cross-attention context—so you can freeze the backbone and still specialize the model for new cell types or tracks.

Key highlights

• Pretrained weights (~1 GB) load via HuggingFace with corrected gamma positions to fix the original port’s rounding errors.
• Accepts DNA sequences as integer indices or one-hot floats, with an optional return_embeddings flag for feature extraction.
• Fine-tuning wrappers include HeadAdapterWrapper, ContextAdapterWrapper, and ContextAttentionAdapterWrapper for new tracks or contextual conditioning.
• GenomicIntervalDataset dynamically pads and lengthens intervals, maps chromosome names between .bed and FASTA files, and returns augmentation metadata.
• Built-in Poisson loss and Pearson correlation coefficient computation for training and evaluation.

Caveats

• The ported checkpoint is explicitly labeled “rough” because numerical differences can still accrue across layers; v0.8.0 fixes the known xlogy mismatch, but the README notes the original port showed absolute errors as high as 0.5, and models instantiated manually via from_hparams still need use_tf_gamma = True to avoid drift.
• Basic training utilities such as gradient accumulation remain on the todo list, so expect to supply your own loop logic.

Verdict

Computational biologists and ML researchers who want Enformer’s long-sequence attention in native PyTorch should start here. If you need a fully packaged training framework with all scaffolding included, it is still a bit too bare-bones.