Ankh: Optimized Protein Language Model Unlocks General-Purpose Modelling

Ahmed Elnaggar; Hazem Essam; Wafaa Salah-Eldin; Walid Moustafa; Mohamed Elkerdawy; Charlotte Rochereau; Burkhard Rost

Source-linked AI summary

Ankh: Optimized Protein Language Model Unlocks General-Purpose Modelling

Ahmed Elnaggar, Hazem Essam, Wafaa Salah-Eldin, Walid Moustafa, Mohamed Elkerdawy, Charlotte Rochereau, Burkhard Rost

arXiv:2301.06568v1Jan 16, 2023cs.LGcs.CLcs.DCq-bio.QM

TL;DR

Protein language modeling has emphasized scaling despite rising computational demands and reduced accessibility. This paper develops Ankh through protein-specific, knowledge-guided optimization and evaluates it across structure, function, and generation tasks. Ankh surpasses the reported state of the art with substantially fewer parameters while generating diverse variants that retain key structural and functional characteristics.

Problem
Protein language models increasingly rely on larger models, but this trend raises computational demands and reduces accessibility.
Method
Ankh uses knowledge-guided experiments across masking, architecture, and pre-training data, with separate High-N and One-N generation frameworks.
Results
Ankh surpassed state-of-the-art performance across the reported benchmarks while using fewer parameters and generated diverse variants retaining key structural and functional characteristics.
Takeaways & Limitations
Ankh demonstrates that protein-specific optimization can support general-purpose modeling with two computation modes and more accessible computational demands.
Takeaways & Limitations
No single formula defined the best model version across the population of candidate formulas, so selection favored holistic generalization across tasks.

Abstract

from arXiv · show

As opposed to scaling-up protein language models (PLMs), we seek improving performance via protein-specific optimization. Although the proportionality between the language model size and the richness of its learned representations is validated, we prioritize accessibility and pursue a path of data-efficient, cost-reduced, and knowledge-guided optimization. Through over twenty experiments ranging from masking, architecture, and pre-training data, we derive insights from protein-specific experimentation into building a model that interprets the language of life, optimally. We present Ankh, the first general-purpose PLM trained on Google's TPU-v4 surpassing the state-of-the-art performance with fewer parameters (<10% for pre-training, <7% for inference, and <30% for the embedding dimension). We provide a representative range of structure and function benchmarks where Ankh excels. We further provide a protein variant generation analysis on High-N and One-N input data scales where Ankh succeeds in learning protein evolutionary conservation-mutation trends and introducing functional diversity while retaining key structural-functional characteristics. We dedicate our work to promoting accessibility to research innovation via attainable resources.

1 Introduction

Protein language models have increasingly pursued scale, but this trend raises computational and accessibility concerns. Ankh instead targets protein-specific optimization and evaluates a general-purpose model across structure, function, and generation settings.

Protein language models have grown from approximately 10^6 to 10^9 parameters, increasing computational demands and reducing accessibility.
Ankh evaluates optimization and generality through structure and function benchmarks plus protein generation on High-N and One-N inputs.
Ankh and Ankh base provide two computation modes designed for different application demands.
The paper frames Ankh as a general-purpose protein language model whose representations are intended to unlock protein sequence-function modeling.

2 Results

Ankh combines knowledge-guided experimentation with compact architectures and evaluates performance, computational efficiency, and protein variant generation. Across the reported analyses, it achieves strong downstream performance while controlling model size and generation diversity.

Computational efficiency and performance: 4.8%: Ankh improved average PLM state-of-the-art performance using less than 10% of training parameters and 30% of embedding dimensions.Ankh base improved performance by 3.4% with 3% of training parameters and 15% of embedding dimensions.
Computational efficiency and performance: 2.2x to 11.7x: ESM-2 (15B) required this range of Ankh base feature-extraction times, while Ankh required 2.0x to 7.1x.The comparison used sequences up to 1024 residues.
High-N generation: The generation analysis reported retention of natural-set entropy patterns while using 500 generated sequences, fewer than 3% of the 16,706 natural sequences.The reported entropy MSE values were 0.1, 0.09, and 0.08 for temperatures 1.0, 1.5, and 2.0.
High-N generation: Temperature 1.0 generated variants as different as 70% identity, while temperature 2.0 reached 55% identity and over 95% of variants were unique.Temperature 1.0 was more conservative, whereas temperature 2.0 produced greater internal variability and broader exploration.
Knowledge-guided optimization: Knowledge-guided experimentation varied masking, architecture, and pre-training data while holding each experiment to two epochs and approximately equal total parameters.The final models integrated top-performing experimental versions across these variable sets.
Model architecture: Ankh’s architecture routes processed sequences through a transformer to residue-level or protein-level prediction, with global max pooling preceding protein-level prediction.Ankh base and Ankh large share this architecture but differ in context-vector length.

3 Discussion

Ankh achieved broad performance gains across protein-modeling benchmarks while using substantially less computational power. Its generation results also preserved evolutionary and functional properties, although model selection involved trade-offs because no single version led every task.

Results: Ankh surpassed state-of-the-art performance across the reported benchmarks, while Ankh base matched or closely approached it with less computational power.The two models provide different computational modes for different application demands.
Results: Ankh learned evolutionary conservation-mutation trends and introduced diversity while retaining key structural and functional characteristics in both High-N and One-N generation.
Interpretation: The findings support protein knowledge-guided model design and the use of sequence contextualized embeddings with task-specific downstream architectures.
Model selection: No single experimental version performed best on every task, so selection favored broad performance, generalization, and negligible numerical differences.The authors note that all experimental versions were trained for only two epochs.
Recommendations: The authors suggest coordinating pre-training data with downstream testing datasets, while identifying that coordination as beyond this work’s scope.

4.1 Self-Supervised Learning

The self-supervised learning setup builds on protein-specific evidence about data choice and transformer architecture. It uses UniRef50 and an encoder-decoder transformer to learn contextual protein representations.

Pre-training datasets: UniRef50 was selected for pre-training because prior ProtTrans analysis associated it with greater sequence variability and representation than UniRef100 and BFD.UniRef50 clusters sequences at a 50% identity threshold, whereas UniRef100 contains identical sequences and sub-fragments.
Transformer architecture: The baseline and experiments use an encoder-decoder transformer that maps input sequences into latent context representations and generates target-domain sequences.
Transformer architecture: ProtTrans experimentation promoted T5 over encoder-only and decoder-only transformer alternatives, motivating the architecture choice used here.

4.2 Downstream Tasks

The downstream evaluation spans protein structure, function, annotation, localization, and engineering-oriented prediction and generation tasks. It combines residue-level and protein-level assessments with family-based and single-sequence variant generation.

Benchmark scope: The benchmark covers regression, classification, and generation across residue-level and protein-level protein-understanding tasks.The evaluation includes structure and function benchmarks alongside protein-engineering generation analyses.
Protein function prediction: FluP predicts green fluorescent protein mutant fluorescence, while SolP classifies proteins as soluble or insoluble.FluP represents genotype-phenotype fitness prediction, and SolP uses identity-filtered dataset splits.
Protein function prediction: GB1 evaluates binding fitness after mutations at four positions, representing non-additive mutation interactions known as epistasis.The task uses the FLIP benchmark and focuses on a restricted four-position mutation landscape.
Protein structure prediction: Structure tasks include residue-pair contact prediction, 1194-class fold prediction, and 3- or 8-class secondary-structure prediction.Contact prediction targets three-dimensional residue contacts, while fold prediction tests generalization to held-out clustered superfamilies.
Annotation and localization: Embedding-based annotation transfer uses nearest neighbors in embedding space to transfer annotations without supervised training or multiple-sequence-alignment calculations.The evaluation uses annotated lookup sequences and query sequences from the ProTucke benchmarking set.
Annotation and localization: The study also evaluates localization into 10 sub-cellular classes and variant generation in family-based and single-sequence settings.High-N generation uses a curated malate dehydrogenase family dataset associated with ProteinGAN.

4.3 Downstream Model: ConvBERT

The downstream models map Ankh embeddings to task-specific targets using a largely shared supervised network, with pooling and configuration adjusted for task requirements.

Task-specific downstream modeling: The supervised downstream models share a network design with minor modifications for residue-level versus protein-level processing and different output distributions.The output types include binary classification, multiclass classification, and regression.
Task-specific downstream modeling: Global max pooling is used only for regression and binary classification tasks before or at the beginning of the shared supervised network.
Task-specific downstream modeling: The authors promote task-specific optimization because different top models and hyperparameters can improve downstream performance.They nevertheless unify the setting of the evaluated models for comparison.

4.4 Variant Generation Model

The variant-generation framework uses different fine-tuning strategies for family-based and single-sequence generation, with temperature and masking controls governing exploration and exploitation.

High-N generation: Auto-regressive fine-tuning generates variants for High-N protein families by freezing the encoder and updating the decoder.The framework is designed to scale across protein families and uses logit-warping temperature to control exploration versus exploitation.
One-Shot generation: Masked Language Modeling is used for One-Shot generation from a single protein sequence.The two generation frameworks are presented separately because their evaluation settings differ from the pretrained-top model framework.
Generation controls: 50% masking with temperature 1.0 and 40% masking represent alternative exploration-exploitation settings for generation.The experiments vary masking probability and logit-warping temperature to explore different generation behaviors.

4.5 Computational Power (Software & Hardware)

The work addresses computational efficiency through software and hardware choices, using JAX and Flax with TPU-v4 infrastructure for neural-network training.

Software: JAX provides just-in-time compilation, automatic differentiation, CPU/GPU/TPU compatibility, and vectorization and parallelization.JAX can deliver approximately 1.4x speed-up for language-model training on TPU Pods compared with PyTorch.
Software: Flax serves as the end-to-end high-performance library and ecosystem used to build the models.It is designed for flexibility and tailored for neural networks within the JAX ecosystem.
Hardware: The pretrained models were trained on Google TPU-v4 Pods with either 64 or 128 cores.A single TPU-v4 VM host contains 8 TPU cores, each with 16 GiB of high-bandwidth memory.

4.6 Data & Model Experimentation

The experiments use protein-specific masking, architecture, and positional-embedding choices to identify configurations that improve performance and support generalization across tasks.

Masking: Masking experiments varied strategy and probability to adapt pre-training to skewed amino-acid distributions and database redundancy.The study tested five masking-strategy variations and additional masking probabilities against a baseline.
Masking: Masking every unique amino acid at least once with full sequence reconstruction outperformed the baseline, whereas 3-gram masking reduced downstream performance.The unique-token strategy was retained for subsequent experiments.
Masking: 1-Gram Span Partial De-masking/Reconstruction merges consecutive unmasked output tokens and achieved higher performance than the first experiment.The change reduces computation associated with reconstructing unmasked output tokens while leaving input tokens unchanged.
Masking: 10% masking was the worst tested probability, while 15% and 30% performed better on different task groups.The authors selected an intermediate value to promote generalization across diverse tasks rather than optimize a single task.
Architecture: The 48-layer encoder and 24-layer decoder configuration outperformed the comparison across secondary structure, fold, localization, and overall mean and median results.The configuration was retained because it supports richer embeddings, generation, generalization, and faster feature extraction.
Architecture: None of the tested depth-width combinations consistently surpassed Exp. 11, so the authors reverted to Exp. 11.This comparison set changed the activation function to classic ReLU while testing alternative depth-width combinations.
Positional embedding: Relative Positional Embedding represents positions through key-query offsets rather than fixed position-specific embeddings.The simplified implementation adds a shared scalar representation to the attention logits across model layers.

5 Availability

Ankh and Ankh base are publicly available for research use through the Ankh repository, which includes tutorials and examples for downstream applications.

Availability: Both Ankh models are available in the Ankh repository under a CC BY-NC-SA 4.0 license.The repository provides access for research innovation.
Availability: The repository includes Python and Jupyter tutorials for embedding extraction and supervised learning on downstream tasks.The examples use freely available online resources such as Google Colab.

7 Supplementary Materials

The supplementary experiments compare pre-training data, masking, architecture, and dataset choices to identify configurations that improve general protein-language-model performance efficiently. The selected design favors broadly consistent performance across tasks rather than a single universally best configuration.

Masking strategy: Exp. 4 was the top masking-strategy version because it reconstructed consecutive unmasked tokens as one merged token.This version proceeded to the masking-probability experiments.
Masking probability: The 30% masking probability led contact prediction and CASP12 secondary-structure results, while 15% performed better on several other tasks.The authors promoted the higher probability for generalization across task types and used CASP12 where secondary-structure results were inconsistent.
Architecture: Exp. 11, with a 48-layer encoder and 24-layer decoder, outperformed Exp. 8 across several secondary-structure, fold, localization, mean, and median evaluations.The choice was motivated by richer embeddings from a larger encoder while retaining decoder layers for generation tasks.
Architecture: Depth-width variations did not consistently surpass Exp. 11, so the authors reverted to Exp. 11 and defined the 768-dimensional version as Ankh base.The tested depth-width configurations produced fluctuating results.
Architecture: Relative positional encoding selected Exp. 20 for its most consistent and general results, whereas weight tying failed to consistently surpass the top-performing model.Exp. 20 used an embedding offset of 128 and 64 embeddings; the weight-tying outcome was attributed to input-output differences under masking and demasking.

7.2 Model Pre-training

The paper pre-trains Ankh models and evaluates their contextualized embeddings across diverse supervised protein tasks, attention-based contact prediction, annotation transfer, and generated-domain analyses. These evaluations cover both predictive benchmarking and whether generated sequences retain or diversify relevant domain characteristics.

Model pre-training: Ankh and Ankh base are fully pre-trained models whose configurations are compared with the baseline and downstream model settings.The study also reports the supervised downstream datasets and their task descriptions.
Embedding-based benchmarking: Embedding-based benchmarking compares Ankh models with ProtT5-XL-U50 and ESM variants using task-specific correlation, accuracy, and contact-precision metrics.The benchmark spans regression, classification, contact prediction, secondary structure, fold, localization, and annotation-transfer tasks.
Contact prediction: Attention-based contact prediction performed worse than embedding-based contact prediction in the reported benchmarking.The attention maps and contextualized embeddings were both tested for fair comparison across models.
CATH analysis: CATH annotation transfer evaluates predicted labels against truth labels across Class, Architecture, Topology, and Homologous Superfamily levels.The analysis reports accuracy over the four CATH levels using the closest match in embedding space.
Generated-domain analysis: Generated and natural sequences are compared through CATH super-family and functional-family domain distributions.The functional analysis distinguishes the target malate dehydrogenase family from the closely similar L-lactate dehydrogenase family.