enformer-pytorch: A PyTorch Port of Enformer With Fine-Tuning Adapters

The gap enformer-pytorch fills for PyTorch genomics teams

DeepMind released Enformer as TensorFlow Sonnet code. The README links to that original implementation and describes this repository as an implementation of the same attention network in PyTorch, with the additional means to fine-tune pretrained models for downstream tasks. The practical problem is not that Enformer is unavailable. It is that the reference implementation lives in a framework that many sequence-modeling teams do not train in. If your data loaders, distributed training setup, and experiment tracking are PyTorch, porting a model by hand is a project in itself. This repository packages that port as a pip install, plus a Hugging Face hosted checkpoint and adapter modules for fine-tuning. The intended audience is a research engineer who wants Enformer as a component rather than as a paper reproduction. The README also notes a finetuned model for predicting pseudobulk chromatin accessibility, pointing at a bioRxiv preprint, which suggests the fine-tuning path has been used beyond the toy examples in the documentation.

What the model returns: heads, embeddings, and the 196,608 base input

The mechanism visible in the README is a single forward pass over a fixed-length sequence. Input is either integer-encoded bases or one-hot floats. The example builds a tensor of shape (1, 196_608) with values from 0 to 5 standing for ACGTN, and notes that -1 is used for padding. The model is constructed through Enformer.from_hparams with explicit dim, depth, heads, output_heads, and target_length arguments. The example configuration is dim 1536, depth 11, heads 8, output_heads as a dict mapping human to 5313 and mouse to 1643, and target_length 896. The forward pass returns a dict keyed by species: output['human'] has shape (1, 896, 5313) and output['mouse'] has shape (1, 896, 1643). So the architecture is multi-head at the species level, and the output is a track-by-position tensor rather than a single prediction. Passing return_embeddings = True changes the return signature to a tuple of (output, embeddings), where embeddings has shape (1, 896, 3072). That 3072 figure is the fine-tuning surface: the adapter wrappers described later attach to this representation rather than to the species heads.

The training loop and the correlation coefficient metric

Training is handled inside the model call rather than in a separate loss function. You pass head and target into forward, and the README states you get the Poisson loss back. The example moves the model to CUDA, builds a target of shape (200, 5313), and calls backward on the returned loss. After training, passing return_corr_coef = True to the same call returns a Pearson R, which the README identifies as the metric used in the paper. That design choice keeps the loss and the evaluation metric in one place, which is convenient, but it also means the training semantics are baked into the module. If you want a different loss, or you want to compute the metric on a separate held-out pass with different batching, you are working around the model's own forward signature. The README does not document an alternative entry point for that. Treat the built-in Poisson loss as the supported path and plan accordingly.

Loading the ported weights and the numerical discrepancy the README admits

The pretrained weights come from DeepMind's TensorFlow Sonnet release, ported and uploaded to Hugging Face as EleutherAI/enformer-official-rough, roughly 1 GB. Loading is one call: from_pretrained('EleutherAI/enformer-official-rough'). The README is unusually candid about the state of that port. It states there are still rounding errors accruing across layers, resulting in an absolute error as high as 0.5, and that the author released a roughly working version because the correlation coefficient looked good. The suspected source is named: the attention pooling module, where attention logits were observed to be high. A later update reports that John St. John found the model hits the paper's reported marks, with human Pearson R of 0.625 on validation and 0.65 on test. A further update says that as of version 0.8.0, from_pretrained automatically uses precomputed gamma positions to address a difference between TensorFlow and PyTorch xlogy, which should resolve the discrepancy. The caveat matters: if you fine-tune without from_pretrained, the README instructs you to set use_tf_gamma = True when instantiating via from_hparams. Miss that flag and you are back in the regime the original release warned about. The README also provides a sanity check, python test_pretrained.py, which it says reports 0.5963 correlation coefficient on a validation sample.

Fine-tuning paths: head, context, and attention adapters

Three wrappers in enformer_pytorch.finetune cover different fine-tuning shapes. HeadAdapterWrapper takes the loaded enformer, a num_tracks argument, and a post_transformer_embed flag. The README explains the flag's meaning directly: by default embeddings are taken after the final pointwise block with conv and gelu, and setting the flag to True takes them right after the transformer block with a learned layernorm. The example uses num_tracks = 128 and a target of shape (1, 200, 128). ContextAdapterWrapper adds a context_dim argument, and the example feeds a context tensor of shape (4, 1024) alongside a target of shape (1, 200, 4), matching four contexts to four tracks. This is the path for cell type or transcription factor conditioning. ContextAttentionAdapterWrapper is described as using attention aggregation from a set of context embeddings, or a single context embedding. For memory-limited fine-tuning, from_pretrained accepts use_checkpointing = True, and target_length can be overridden at load time, shown with target_length = 128 and dropout_rate = 0.1 for shorter sequences. That override is the escape hatch when the full 196,608 base input does not fit your budget.

Where this port is the wrong tool

The admitted numerical history is the first limitation, and it is not fully closed by a version note. The README says the gamma fix should resolve the discrepancy, which is weaker than a statement that it has been verified across the full output. If your work depends on absolute predicted values rather than rank correlation, the documented absolute error of up to 0.5 in the earlier release is a reason to validate against the TensorFlow reference on your own tracks before trusting outputs. Second, the fine-tuning API assumes a particular shape of problem: sequence in, tracks or contexts out, Poisson loss. If your downstream task is classification, or your labels are not counts, the built-in loss is a mismatch and the README does not describe how to substitute one. Third, the model is large by construction. The example configuration is dim 1536 with depth 11, and the checkpoint is about 1 GB, which is why use_checkpointing exists. Fourth, the documentation is example-driven. There is no described CLI, no config file format, and no described preprocessing pipeline for turning raw genomic intervals into the 196,608-length input. You supply that.

The TensorFlow Sonnet implementation is the real alternative

The obvious alternative is the original DeepMind code, which the README links to directly in the deepmind-research repository. The difference in approach is not cosmetic. The TensorFlow version is the reference the PyTorch port was validated against, so it is the thing you compare to when you suspect a numerical problem. It also carries no porting risk, because there is no port. The cost is that you inherit a TensorFlow and Sonnet stack, and the fine-tuning adapters described here, HeadAdapterWrapper, ContextAdapterWrapper, and ContextAttentionAdapterWrapper, have no counterpart there. So the choice is between framework fit and reference fidelity. If your goal is to reproduce reported metrics or to audit a discrepancy, use the TensorFlow code. If your goal is to train Enformer on new tracks inside an existing PyTorch codebase, use this repository and accept that you are relying on a port whose numerical correctness the README describes as an ongoing concern that a gamma-position fix is expected to address.

Maintenance, releases, and the MIT licence

The repository is not archived and the most recent push is dated 2026-06-26. Releases are sparse: 0.8.11 in July 2025, then 0.8.10 and 0.8.9 in October 2024. The README's own version guidance is looser than the release list, telling users to install enformer-pytorch>=0.5 for the gamma behaviour and stating that the fix landed in 0.8.0. That gap between the documented minimum and the actual latest release is worth noting when you pin a version. The project is MIT licensed, which is permissive and places few constraints on commercial or derivative use. Two caveats that are not legal advice: the MIT licence covers this repository's code, and the README does not state the licence terms of the ported weights hosted on Hugging Face, which originate from DeepMind's release, so check those separately before redistribution. Also check the terms attached to any finetuned checkpoint you pull, since those are separate artifacts from the code.

Editorial conclusion

Adopt enformer-pytorch if you are already in PyTorch and want to fine-tune Enformer on your own tracks or contexts through HeadAdapterWrapper, ContextAdapterWrapper, or ContextAttentionAdapterWrapper. Do not adopt it if you need a bit-exact reference for reproducing the paper, or if you cannot tolerate the documented rounding errors in the ported weights. Before committing, run test_pretrained.py against EleutherAI/enformer-official-rough and confirm the correlation coefficient you get on your own validation sample, then check whether you need use_tf_gamma = True when instantiating from hparams rather than from_pretrained.

Official sources

1. Issues
2. License: MIT
3. lucidrains/enformer-pytorch on GitHub
4. README
5. Releases