Modern Hopfield Networks and Attention for Immune Repertoire Classification

Michael Widrich; Bernhard Schäfl; Hubert Ramsauer; Milena Pavlović; Lukas Gruber; Markus Holzleitner; Johannes Brandstetter; Geir Kjetil Sandve; Victor Greiff; Sepp Hochreiter; Günter Klambauer

Source-linked AI summary

Modern Hopfield Networks and Attention for Immune Repertoire Classification

Michael Widrich, Bernhard Schäfl, Hubert Ramsauer, Milena Pavlović, Lukas Gruber, Markus Holzleitner, Johannes Brandstetter, Geir Kjetil Sandve, Victor Greiff, Sepp Hochreiter, Günter Klambauer

arXiv:2007.13505v1Jul 16, 2020cs.LGq-bio.BMstat.ML

TL;DR

The paper addresses immune repertoire classification as massive MIL, where disease-relevant sequences are rare within very large bags. It introduces DeepRC, integrating transformer-like attention with modern Hopfield networks, and reports superior predictive performance across simulated and real-world experiments while extracting disease-associated motifs.

Problem
Immune repertoire classification requires learning from hundreds of thousands of sequences per bag without instance-level labels and with witness rates as low as 0.01%.
Method
DeepRC uses transformer-like attention, equivalently a modern Hopfield-network update rule, to pool learned representations of immune-receptor sequences for repertoire-level classification.
Results
DeepRC outperforms the compared methods across simulated and real-world experimental conditions and extracts sequence motifs connected to disease classes.
Takeaways & Limitations
Modern Hopfield attention provides a practical mechanism for identifying discriminating sequences and motifs in immune repertoires containing massive numbers of instances.
Takeaways & Limitations
Most repertoire sequences lack known antigen assignments, limiting direct experimental verification and interpretation of machine-learning-identified motifs.

Abstract

from arXiv · show

A central mechanism in machine learning is to identify, store, and recognize patterns. How to learn, access, and retrieve such patterns is crucial in Hopfield networks and the more recent transformer architectures. We show that the attention mechanism of transformer architectures is actually the update rule of modern Hopfield networks that can store exponentially many patterns. We exploit this high storage capacity of modern Hopfield networks to solve a challenging multiple instance learning (MIL) problem in computational biology: immune repertoire classification. Accurate and interpretable machine learning methods solving this problem could pave the way towards new vaccines and therapies, which is currently a very relevant research topic intensified by the COVID-19 crisis. Immune repertoire classification based on the vast number of immunosequences of an individual is a MIL problem with an unprecedentedly massive number of instances, two orders of magnitude larger than currently considered problems, and with an extremely low witness rate. In this work, we present our novel method DeepRC that integrates transformer-like attention, or equivalently modern Hopfield networks, into deep learning architectures for massive MIL such as immune repertoire classification. We demonstrate that DeepRC outperforms all other methods with respect to predictive performance on large-scale experiments, including simulated and real-world virus infection data, and enables the extraction of sequence motifs that are connected to a given disease class. Source code and datasets: https://github.com/ml-jku/DeepRC

Introduction

The paper connects transformer attention to modern Hopfield-network updates and applies this equivalence to immune repertoire classification, a massive multiple instance learning problem.

Problem setting: Existing MIL methods typically handle hundreds or a few thousand instances per bag, whereas immune repertoires contain hundreds of thousands with witness rates down to 0.01%.
Hopfield networks and attention: Transformer attention is identified as the update rule of continuous modern Hopfield networks, whose storage capacity grows exponentially with representation dimension.This capacity motivates using modern Hopfield networks to extract patterns from very large instance sets.
Problem setting: Immune repertoire classification is a MIL problem in which a disease-related immune status must be inferred from immense bags of receptor sequences.Only a small fraction of receptors may determine status, while sequence overlap is low and pathogen binding can depend on subsequences.
DeepRC: DeepRC applies transformer-like attention to sequence-level MIL, pooling learned sequence representations into repertoire representations rather than classifying individual sequences.The method is designed for large bags of complex immune-receptor sequences.
DeepRC: DeepRC addresses limitations of prior approaches by using attention pooling, representation pooling, and flexible feature extractors such as one-dimensional convolutions or LSTMs.

Exponential storage capacity of continuous state modern Hopfield networks with transformer attention as update rule

Modern Hopfield networks use an update rule identical to transformer attention and can store exponentially many patterns, supporting retrieval from massive instance sets. Their storage capacity motivates applying the mechanism to massive multiple-instance learning.

Transformer attention equivalence: The update rule ξnew = f(ξ; X, β) = X softmax(βX^Tξ) is also the formula of transformer attention.The update globally converges to stationary points of the energy function.
Retrieval behavior: When patterns are well separated, one update typically retrieves a nearby pattern with exponentially small retrieval error.Poorly separated patterns instead converge near their arithmetic mean, while similar clusters can create metastable states.
Definitions: The modern Hopfield formulation defines stored patterns through isolated spheres and retrieved patterns through convergence to a single fixed point.These definitions formalize storage and retrieval before the capacity analysis.
Storage capacity: For randomly chosen patterns, modern Hopfield networks can store a number of patterns exponential in the pattern-space dimension d.The theorem gives probability-controlled storage bounds under specified radius, inverse-temperature, and separation assumptions.
Motivation for massive MIL: The attention mechanism’s storage capacity is critical for massive MIL because it can retrieve a few relevant instances from bags containing very many patterns.The paper explicitly connects exponential storage and retrieval to MIL with massive numbers of instances.

Deep Repertoire Classification

DeepRC formulates immune repertoire classification as massive multiple instance learning and uses transformer-like attention to pool sequence representations into an interpretable repertoire representation. Its attention mechanism is linked to modern Hopfield networks and assigns sequence-level importance for classification.

Problem setting: DeepRC treats each immune repertoire as an unordered bag of immune-receptor sequences with variable size, where only bag labels are observed and instance labels remain hidden.This creates a credit-assignment problem because the sequences responsible for the repertoire label must be identified.
Architecture: DeepRC applies transformer-like key-value attention to pool learned sequence representations rather than instance-level predictions, followed by an output network that predicts immune status.The model uses a convolutional sequence encoder for values, a second network for keys, and a fully connected classifier for the repertoire representation.
Scale: DeepRC is designed for immune repertoires containing very large numbers of sequences, extending attention-based MIL beyond the hundreds or thousands of instances typical in earlier applications.The method targets the large-instance, low-witness-rate setting posed by immune repertoire classification.
Attention pooling: The attention mechanism produces a fixed-size repertoire representation as a weighted mean of sequence representations, enabling the model to process bags with different numbers of instances.The attention weights are computed over the sequence representations and determine their contribution to the pooled vector.
Interpretability: Attention weights provide an interpretable estimate of each sequence’s importance, and DeepRC also supports contribution-analysis methods such as Integrated Gradients.These outputs can be used to inspect which sequences contribute to the repertoire-level prediction.
Training and inference: DeepRC is trained end to end with gradient descent and uses an implementation that retains the top 10% of sequences by attention weight for updates and prediction.The reported training setup uses cross-entropy loss, Adam optimization, and sequence dropout.

Experimental Results

DeepRC is evaluated against multiple baselines across simulated, generated, implanted-signal, and real-world immune repertoire datasets using ROC-AUC. It outperforms competing methods across dataset categories, while performance becomes difficult at extremely low witness rates.

Datasets: 30 billion sequences across four dataset categories support comparisons spanning simulated, generated, implanted-signal, and real-world immune repertoire classification.Repertoires contain approximately 300,000 sequences on average, except low-coverage real-world data with 10,000 sequences per repertoire.
Experimental setup: The study uses nested 5-fold cross-validation, with methods tuning important hyperparameters on inner validation sets.Table 1 reports standard deviations across five cross-validation folds except for the simulated column.
Results: DeepRC outperforms all competing methods in average AUC across real-world, implanted-signal, LSTM-generated, and simulated dataset categories.Across categories, the runner-up is either the SVM with MinMax kernel or the burden test.
Results on simulated immunosequencing data: DeepRC reaches 0.846 ± 0.223 average AUC on simulated data, versus 0.827 ± 0.210 for the SVM with MinMax kernel.All methods score below 0.550 when only 0.01% of sequences carry the motif.
Results on real-world data with implanted motifs: DeepRC achieves 0.980 ± 0.029 average AUC on real-world data with implanted motifs, compared with 0.883 ± 0.170 for the burden test.Other methods struggle with noisy motifs at a frequency of 0.1%.

Broader Impact

The paper discusses possible scientific, diagnostic, and societal applications of DeepRC alongside limitations involving motif validation, wet-lab follow-up, and dataset biases. It emphasizes that successful deployment would require experimental verification and attention to confounding factors.

Scientific impact and limitations: Machine-learning-identified motifs often lack known corresponding antigens, making experimental verification and interpretation an important remaining challenge.High-throughput antigen-labeled immunosequencing may help address this limitation.
Impact on society: DeepRC could support faster immune-status testing from blood samples if the approach proves successful, potentially affecting diagnostic and disease-tracking pipelines.The paper presents these as possible downstream applications rather than established outcomes.
Impact on society: Successful identification of resistance-associated insights could eventually support faster targeted immunisation, but this remains conditional on the approach proving successful.The paper frames this as a possible societal consequence rather than a demonstrated result.
Consequences of failures: Applying DeepRC without reflection could produce failures that are detected through unsuccessful wet-lab validation, while the algorithm itself does not directly prescribe treatment.The proposed pipeline includes wet-lab tests after algorithmic analysis.
Bias and discrimination: Age, sex, and other confounding factors could bias predictions or produce uneven performance across subgroups.Available data on confounders may improve performance for some groups more than others.

Appendix

The appendix provides implementation, dataset, experimental, interpretability, motif-extraction, and ablation details supporting the main paper.

Appendix: The appendix documents DeepRC implementation details, compared methods, experimental setup and results, dataset generation, interpretation, motif extraction, and ablation studies.

1 Immune Repertoire Classification

The supplied passage contains only the appendix table of contents, including an introduction section and related implementation and dataset sections.

1 Immune Repertoire Classification: The appendix contents list an introduction, DeepRC implementation details, and dataset sections including simulated and LSTM-generated data.

Immune Repertoire Classification

The supplied passage lists the paper’s authors.

The author list includes Michael Widrich, Bernhard Schäfl, and Milena Pavlović.
The listed authors also include Geir Kjetil Sandve, Sepp Hochreiter, and Victor Greiff.
Günter Klambauer is included among the paper’s listed authors.

A1 Introduction

The appendix is organized into sections covering architecture, datasets, comparisons, hyperparameters, results, sequence generation, interpretation, and related details.

Section A2 details the DeepRC architecture, while Section A3 describes the datasets.
Sections A4 and A5 explain the compared methods and hyperparameter-selection process.
Section A6 presents detailed dataset-category results in tabular form.

A2 DeepRC implementation details

DeepRC combines amino-acid sequence encoding, convolutional motif extraction, attention-based aggregation, subsampling, and computationally constrained training, while remaining limited by computational cost and over-fitting risk.

Input layer: DeepRC concatenates 20-dimensional amino-acid one-hot vectors with three positional features encoding sequence beginning, center, and end.The positional features describe each amino acid’s relative location in the sequence.
Motif recognition: A lightweight 1D CNN extracts sequence motifs, with kernel size tuned during hyperparameter search and alternative RNN or transformer encoders discussed.The CNN produces fixed-size sequence representations from learned local patterns.
Regularization and computation: Training subsamples each repertoire to 10,000 sequences and uses mixed-precision computation, small batches, and GPU-oriented implementation measures.These choices reduce computational effort and GPU-memory demands.
Limitations: DeepRC’s main limitations are computationally demanding hyperparameter and model selection, GPU-memory requirements, and over-fitting with relatively few real-world bags.More computational power and GPU memory could support larger or more advanced models.

A3 Datasets

The benchmark uses simulated and real immune receptor repertoires with implanted sequence motifs across dataset categories varying in realism and signal difficulty.

The study constructs classification scenarios with varying realism and difficulty to compare machine-learning methods.
Positive-class repertoires receive implanted sequence motifs that simulate immune signals through short biologically meaningful sequence stretches.The motivation is that immune-receptor interactions with antigens occur through short sequence stretches.
The benchmark includes simulated data, LSTM-generated data, real-world data with implanted signals, and real-world immunosequencing data.
Within the first three categories, datasets vary in difficulty according to implanted-signal type and the proportion of sequences carrying the signal.

A3.1 Simulated immunosequencing data

The study creates 18 simulated immunosequencing datasets to test how signal frequency and complexity affect classification when repertoires contain many random amino-acid sequences.

A3.1 Simulated immunosequencing data: 18 datasets vary the frequency and complexity of implanted signals to probe their effects on method performance.Each dataset contains many repertoires with many random amino-acid sequences, and signals are implanted in positive-class repertoires.
A3.1 Simulated immunosequencing data: The simulations reproduce real-world CMV repertoire distributions for amino-acid frequencies, sequence lengths, and instances per bag.
A3.1 Simulated immunosequencing data: The resulting 18 datasets contain roughly 28.7 billion sequences and span variable difficulty through different implanted-motif properties.

A3.2 LSTM-generated data

The LSTM-generated-data experiments use autoregressive sequence generation to create immune-receptor-like repertoires, then implant noisy motifs with varied witness rates for classification tests.

A3.2 LSTM-generated data: An LSTM trained on CMV immunosequences generates repertoires with sequence properties resembling experimentally observed immune receptors.The model captures positional distributions, amino-acid combinations, relationships, and ordering before autoregressive generation.
A3.2 LSTM-generated data: Five datasets vary motif witness rates and properties by implanting noisy length-4 motifs into positive repertoires.Each positive-class sequence carries the motif with probability ρ, and motif positions are probabilistically altered during implantation.
A3.2 LSTM-generated data: Additional real-sequence experiments implant altered motifs at rates of 1% or 0.1% across four low-coverage dataset variations.These datasets contain 750 repertoires per class, with 10k sequences per repertoire, and use experimentally obtained immunosequences.
A3.2 LSTM-generated data: The generated sequences exhibit distributions of 4-mers and amino acids similar to the original CMV dataset.
A3.2 LSTM-generated data: DeepRC extracts implanted motifs from simulated models and identifies complex, highly variable central patterns in the real-world CMV dataset.