Enformer

Transformer-based deep neural network for predicting genomic coverage tracks from long DNA sequences with long-range context.

Disclaimer

This is an UNOFFICIAL implementation of Effective gene expression prediction from sequence by integrating long-range interactions by Žiga Avsec, Vikram Agarwal, Daniel Visentin, et al.

The OFFICIAL repository of Enformer is at google-deepmind/deepmind-research/enformer.

The MultiMolecule team has confirmed that the provided model and checkpoints are producing the same intermediate representations as the original implementation.

The team releasing Enformer did not write this model card for this model so this model card has been written by the MultiMolecule team.

Model Details

Enformer is the successor of Basenji. It replaces Basenji's dilated convolution tower with a convolution stem followed by a Transformer trunk, which lets it model long-range genomic interactions. It consumes a long DNA window (~197 kb), passes it through a convolution + attention-pooling stem that downsamples the sequence by 2 ** 7 = 128x, processes the binned representation with 11 Transformer blocks using Transformer-XL style relative positional encoding, center-crops to 896 output bins, and applies a pointwise head plus a per-species linear track projection with a softplus activation. The prediction is binned: the output has shape (batch_size, target_length, num_tracks) where each bin summarizes 128 bp of sequence and num_tracks is the number of genomic coverage experiments for the selected species.

Model Specification

Input Length	Bin Size	Output Bins	Hidden Size	Layers	Heads	Num Labels	Num Parameters (M)	FLOPs (P)	MACs (P)	Max Num Tokens
196608	128	896	1536	11	8	5313	246.18	-	-	196,608

The table reports the human output head. The mouse head predicts 1643 tracks. FLOPs and MACs have not been recomputed for the canonical 196,608 bp Enformer input window.

Links

• Code: multimolecule.enformer
• Data: ENCODE, FANTOM5, GTEx CAGE, ChIP-seq, DNase-seq, and related genomic coverage tracks
• Paper: Effective gene expression prediction from sequence by integrating long-range interactions
• Developed by: Žiga Avsec, Vikram Agarwal, Daniel Visentin, Joseph R. Ledsam, Agnieszka Grabska-Barwinska, Kyle R. Taylor, Yannis Assael, John Jumper, Pushmeet Kohli, David R. Kelley
• Model type: Convolutional stem followed by Transformer trunk with long-range attention for binned multi-track genomic coverage prediction
• Original Repository: google-deepmind/deepmind-research/enformer

Usage

The model file depends on the multimolecule library. You can install it using pip:

pip install multimolecule

Direct Use

Genomic Coverage Prediction

You can use this model to predict binned genomic coverage tracks from a DNA sequence:

>>> import torch
>>> from multimolecule import DnaTokenizer, EnformerConfig, EnformerForTokenPrediction

>>> config = EnformerConfig(
...     sequence_length=256, hidden_size=12, num_hidden_layers=1, num_attention_heads=2,
...     attention_head_size=4, num_downsamples=3, dim_divisible_by=2, target_length=16,
...     num_labels=4,
... )
>>> model = EnformerForTokenPrediction(config)
>>> output = model(torch.randint(config.vocab_size, (1, 256)))
>>> output.logits.shape
torch.Size([1, 16, 4])
>>> coverage, channels = model.postprocess(output)
>>> coverage.shape
torch.Size([1, 16, 4])

The binned positional axis is treated as the "token" axis: each output position corresponds to one genomic bin rather than a single nucleotide. The species configuration option selects the human (5,313 tracks) or mouse (1,643 tracks) output head.

Interface

• Input length: fixed 196,608 bp DNA window
• Output binning: 128 bp per output bin; 896 output bins per window (after center-cropping the binned representation)
• Species head: select human (5,313 tracks) or mouse (1,643 tracks) via the species config option
• Output: raw pre-softplus logits of shape (batch_size, target_length, num_tracks); use postprocess for non-negative coverage tracks

Training Details

Enformer was trained to predict genomic coverage tracks (DNase-seq, ATAC-seq, ChIP-seq and CAGE) from the human and mouse reference genomes.

Training Data

The model was trained on a large compendium of functional genomics experiments aligned to the human (hg38) and mouse (mm10) reference genomes. The genome was divided into overlapping windows; for each window the per-128-bp coverage of every experiment served as the regression target.

Training Procedure

Pre-training

The model was trained to minimize a Poisson regression loss between predicted and observed coverage, using a softplus output activation to keep the predicted coverage non-negative.

Citation

@article{avsec2021effective,
  author    = {Avsec, {\v{Z}}iga and Agarwal, Vikram and Visentin, Daniel and Ledsam, Joseph R. and Grabska-Barwinska, Agnieszka and Taylor, Kyle R. and Assael, Yannis and Jumper, John and Kohli, Pushmeet and Kelley, David R.},
  title     = {Effective gene expression prediction from sequence by integrating long-range interactions},
  journal   = {Nature Methods},
  year      = 2021,
  volume    = 18,
  number    = 10,
  pages     = {1196--1203},
  doi       = {10.1038/s41592-021-01252-x},
  publisher = {Nature Publishing Group}
}

The artifacts distributed in this repository are part of the MultiMolecule project. If MultiMolecule supports your research, please cite the MultiMolecule project as follows:

@software{chen_2024_12638419,
  author    = {Chen, Zhiyuan and Zhu, Sophia Y.},
  title     = {MultiMolecule},
  doi       = {10.5281/zenodo.12638419},
  publisher = {Zenodo},
  url       = {https://doi.org/10.5281/zenodo.12638419},
  year      = 2024,
  month     = may,
  day       = 4
}

Contact

Please use GitHub issues of MultiMolecule for any questions or comments on the model card.

Please contact the authors of the Enformer paper for questions or comments on the paper/model.

License

This model implementation is licensed under the GNU Affero General Public License.

For additional terms and clarifications, please refer to our License FAQ.

SPDX-License-Identifier: AGPL-3.0-or-later