Attention patterns inside a protein language model used to cut sequences into reusable units

The work

Technical · from the record

The study parses protein sequences into units it calls 'protein words' by binarizing the 660 attention matrices that the ESM2-650M protein language model produces for a sequence and clustering residues with the Louvain community detection algorithm, then compiles dictionaries of high-occurrence words from 1 million UniRef50 sequences and from individual Pfam families (20 762 family dictionaries). A supervised transformer model, Word2Function, was trained on word embeddings from 56 299 sequences labelled with 65 GO molecular-function terms, and Integrated Gradients attributions were used to assign individual words to GO terms in a table called WordTableGO65. On 182 proteins from the ProteinGym deep mutational scanning set, median functional residue coverage was 0.900 with the combined Pfam and UniRef50 dictionaries, compared with 0.045 for the motif-based tool PROSITE and 0.017 for InterProScan; for whole-protein GO-term prediction on the PWNet dataset, mean functional MCC was 0.284 for Word2Function and 0.150 for PROSITE. The word table was also matched against 1 494 146 InterPro sequences annotated as domains of unknown function and against MHC peptides from the MHC Motif Atlas, where the UniRef50 dictionary matched 2832 MHC-I peptides versus 373 for a random k-mer dictionary.

How AI was used

A pre-trained protein language model, ESM2-650M, was run without fine-tuning on each input sequence (length up to 1024, canonical residues only) to obtain 660 attention matrices and 1280-dimensional per-residue embeddings. Each attention matrix was binarized with a dual cutoff (a fixed threshold of 0.1, plus a proportional threshold discarding the lowest 20% of scores in heads with diffuse attention), converted to a directed residue graph, and partitioned with the Louvain algorithm (resolution 1.0, modularity gain threshold 1e-7); communities of 5–20 residues, contiguous or with up to two gaps, were kept as raw words. Running this in multi-sequence mode over 1 million UniRef50 sequences (50 per each of 20 000 Pfam families) and over individual families produced a common dictionary and family dictionaries, with the 20 amino acids collapsed into 12 degenerate types and length-stratified occurrence thresholds; in single-sequence mode, raw words for an analyte sequence were degenerated and matched exactly against these dictionaries to prioritize words. For function mapping, word embeddings were formed by averaging the ESM2 residue embeddings inside each word and fed to transformer layers plus a fully connected layer (embedding and hidden dimension 1280, 20 attention heads) trained with binary cross-entropy, Adam, batch size 256, weight decay 1e-5, initial learning rate 1e-3, cosine decay and early stopping, to predict 65 GO terms over the ExpGO65 sequences. Integrated Gradients was then applied to the trained model to score each word's contribution, and words whose cumulative positive attribution reached a cutoff of half the total were written to the WordTableGO65 table, which was subsequently used for whole-protein annotation by exact word matching. Baselines were obtained by submitting sequences to ScanProsite and InterProScan, retraining ProtNote on the same ExpGO65 splits, and querying the GPSFun web server.