Synthyra / ANKH3_xl

FastANKH

Fast, optimized implementations of ANKH protein language models (T5-based) with multi-backend attention support.

Requires PyTorch 2.11+ for Flash Attention 4 (FA4) backend support via flex attention.

Models

Usage

from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("Synthyra/ANKH_base", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("Synthyra/ANKH_base")

# Set attention backend before loading for best performance
# Options: "sdpa" (default, exact), "flex" (FA4 on Hopper/Blackwell)
model.config.attn_backend = "flex"

sequences = ["MKTLLILAVL", "ACDEFGHIKLMNPQRSTVWY"]
inputs = tokenizer(sequences, return_tensors="pt", padding=True)
outputs = model(**inputs)

# Per-residue embeddings
embeddings = outputs.last_hidden_state  # (batch, seq_len, hidden_dim)

Batch Embedding

model = AutoModel.from_pretrained("Synthyra/ANKH_base", trust_remote_code=True).to("cuda")
embeddings = model.embed_dataset(
    sequences=["MKTLLILAVL", "ACDEFGHIKLMNPQRSTVWY"],
    batch_size=8,
    max_len=512,
    full_embeddings=True,
)

Attention Backends

Architecture

ANKH models are T5 encoder-only architectures:

• No absolute position embeddings : Uses T5-style relative position bias (log-bucketed, bidirectional)
• RMS LayerNorm : No mean subtraction, no bias term
• Gated FFN : activation(wi_0(x)) * wi_1(x) -> wo(x) with gelu_new (v1) or silu (v2/v3)
• Pre-layer normalization : Norm before attention and FFN, residual after
• No bias in projections : All q/k/v/o and FFN linear layers are bias=False

The relative position bias is computed once (materialized as a full tensor) and shared across all encoder layers. For the flex backend, the bias is passed as a score_mod closure for optimal performance.

Notes

• The FastAnkhForMaskedLM variant includes an LM head initialized from the shared embedding weights. The original ANKH models were trained with T5's span corruption objective using an encoder-decoder architecture. This encoder-only MaskedLM head is not pre-trained for standard MLM and requires additional fine-tuning.
• ANKH3 models use a vocabulary of 256 tokens (vs 144 for v1/v2) and were trained with dual objectives ([NLU] for embeddings, [S2S] for generation).

Citations

@article{elnaggar2023ankh,
  title={Ankh: Optimized Protein Language Model Unlocks General-Purpose Modelling},
  author={Elnaggar, Ahmed and Essam, Hazem and Salah-Eldin, Wafaa and Moustafa, Walid and Elkerdawy, Mohamed and Rochereau, Charlotte and Rost, Burkhard},
  journal={arXiv preprint arXiv:2301.06568},
  year={2023}
}

@article{alsamkary2025ankh3,
  title={Ankh3: Multi-Task Pretraining with Sequence Denoising and Completion Enhances Protein Representations},
  author={Alsamkary, Hazem and Elshaffei, Mohamed and Elkerdawy, Mohamed and Elnaggar, Ahmed},
  journal={arXiv preprint arXiv:2505.20052},
  year={2025}
}

@misc{FastPLMs,
  author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.},
  title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.},
  year={2024},
  url={https://huggingface.co/Synthyra/ESMplusplus_small},
  DOI={10.57967/hf/3726},
  publisher={Hugging Face}
}

@article{dong2024flexattention,
  title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels},
  author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace},
  journal={arXiv preprint arXiv:2412.05496},
  year={2024}
}

@inproceedings{paszke2019pytorch,
  title={PyTorch: An Imperative Style, High-Performance Deep Learning Library},
  author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith},
  booktitle={Advances in Neural Information Processing Systems 32},
  year={2019}
}