Improving antibody language models using naturally paired data

Antibody language models are computational models that utilize large-scale machine learning techniques, especially self-supervised learning methods, to understand and predict the characteristics and functions of antibody sequences. These models attempt to learn patterns and associations in antibody sequence data by analyzing the amino acid sequence of the antibody, similar to the application of natural language processing techniques to text data. This study used paired and random antibody sequence data to train a language model and found that the model trained using paired data can learn important features of the antibody and perform better on a variety of performance indicators.

A close-up of a white background

Description automatically generated

The human body's antibody library is extremely rich, containing up to 1,018 unique antibodies, far exceeding the total number of proteins in all species on earth. These antibodies are formed through somatic genetic recombination of B cells, involving a combination of different gene fragments. In this study, the authors used the recently released approximately 1.6 million paired antibody sequence data to train antibody-specific language models (AbLMs) in order to explore whether using paired data can significantly improve model performance compared to unpaired data. The research team developed three baseline antibody language model (BALM) variants: BALM-paired (using paired sequences), BALM-shuffled (using randomly paired sequences), and BALM-unpaired (using unpaired sequences). BALM-paired performed significantly better than the other two models on multiple performance indicators, and also demonstrated better performance than the unpaired model on three antibody-specific classification tasks. With these findings, the study reveals the great potential of using paired antibody sequence data to train AbLMs in understanding and predicting antibody properties.

Model architecture

BALM-paired, BALM-shuffled and BALM-unpaired use a slightly modified RoBERTa-large architecture. An encoder-only architecture is employed in order to generate useful sequence embeddings that can be used for downstream tasks such as specificity classification and aligned with existing protein and antibody language models (LMs). When it comes to model training, the only notable exception to codec training is ProtT5. The model was trained using a masked language model (MLM) objective on the same Jaffe dataset, which contains 1,335,854 pairs of antibody sequences. BALM-paired trains on the original naturally paired sequences, BALM-shuffled trains on randomly paired heavy and light chain sequences, and BALM-unpaired trains on separate heavy and light chains (each input has only a chain). To balance training among models, BALM-unpaired has a batch size of 512, which is twice that of BALM-paired and BALM-shuffled, which have a batch size of 256.

Prioritize learning germinal center antibody signatures

A graph of a graph of a loss

Description automatically generated with medium confidence

To evaluate BALM-paired, position-wise cross-entropy loss (CEL) for mutated or unmutated sequences was analyzed separately (fig above). By grouping sequence positions into corresponding framework regions (FRs) or complementarity determining regions (CDRs), the authors observed significantly weaker model performance in non-template CDR3s of unmutated sequences. Additionally, the authors also observed a slight decrease in model performance in all regions of variant antibody sequences containing non-template somatic mutations spread throughout the sequence. Antibody mutations cluster in CDRs, and BALM-paired performs significantly worse in CDRs of variant sequences than FRs with relatively few mutations.

Pairing data prioritizes light chain improvements

A collage of different colored chains

Description automatically generated

Using a test dataset containing 20,000 pairs of naturally paired antibody sequences, the authors analyzed the output embeddings of BALM-paired, BALM-shuffled and BALM-unpaired. As described previously, the input length-dimensional embeddings are averaged from the final Transformer layer and a Uniform Manifold Approximation and Projection (UMAP) representation is computed. Since the output embeddings of BALM-paired and BALM-shuffled include both heavy and light chains, a subset of the embeddings containing only a single chain (heavy or light) was extracted and then averaged in the length dimension. This enables a direct comparison of the embeddings produced by pairwise models and BALM-unpaired. The heavy chain embedding clustering of BALM-paired, BALM-shuffled and BALM-unpaired was mainly grouped by mutations and secondarily by V genes (above fig A-F). However, the light chain embeddings of these models showed quite different clustering (above fig G–L). Although the BALM-unpaired mutated light chain sequences formed relatively clear clusters of V genes, the unmutated light chain insertions were essentially randomly dispersed (above fig K and L). In line with the pattern observed for BALM-shuffled, unmutated light chain embeddings also appeared to be randomly dispersed (above fig I and J). In contrast, BALM-paired produced clustered light chain embeddings similar to those of the heavy chain, grouping primarily by mutations and secondary by V genes (above fig G and H). Given that this improvement in light chain clustering only occurs in BALM-paired and not in BALM-shuffled, this suggests that BALM-paired is learning cross-link features that are only present in naturally paired sequences and that these features preferentially improve light chains Embed.

Pairwise model learns cross-chain features

A diagram of a light chain

Description automatically generated with medium confidence

Next, the authors further delve into the ability of BALM-paired to learn features spanning two antibody chains. From the test data set, the authors selected all sequence pairs containing at least 3 mutations in each heavy and light chain. All mutated heavy chain positions are masked and then BALM-paired predictions are requested for masked residues when the heavy chain is paired with (1) a naturally paired light chain, or (2) a germinal center regression version of the light chain. All mutated light chain residues returned to the germinal center (above fig A). For comparison, BALM-unpaired predictions were also asked for the same masked residues, but only the unpaired heavy chain sequence was given. Because only the mutation position is masked, the prediction of germinal center-encoding residues is always incorrect. For BALM-paired, the authors noticed that the cross-entropy loss (CEL) was significantly reduced when the masked heavy chain was paired with the native (mutated) light chain (above fig B), indicating that native pairing improves model performance through cross-chain learning. BALM-paired considered the correct (mutated) residue about 4 times more likely when paired with the native light chain (6.7 vs. 1.7), and considered incorrect when paired with the native (mutated) light chain but Residues that are not germinal center are approximately twice as likely (16.7 vs. 8.9), indicating that the model is learning patterns of somatic mutations rather than memorizing specific mutations (above fig C).

The results in experiments in which light chain mutations were masked and paired with native or germinal center-reverted heavy chains (above fig D and E) are even more striking: native pairing increases the likelihood of the correct (mutated) residue by more than 6 times (7.3 vs. 1.1) and almost 5 times more likely to be at any non-germinal center residue (11.9 vs. 3.1). To verify that natural pairing, and not just the presence of any paired chains during training, was responsible for the improved performance of BALM-paired, the authors conducted the same experiment using BALM-shuffled. When mutation-masked sequences were paired with mutated or germinal center-regressed partner chains, BALM-shuffled performed indistinguishably from BALM-unpaired, demonstrating that the features learned by BALM-paired are indeed specific to naturally paired antibody sequences.

Question: Would like to discuss what specific ML model have you used in your research work?

Reference: Burbach S M, Briney B. Improving antibody language models with native pairing.

DOI:https://doi.org/10.1016/j.patter.2024.100967

Preprint

Preprint

Improving antibody language models using naturally paired data