Hopfield Pooling: Learned Queries over Sets

Attention Is a Hopfield Update showed that self-attention is one step of a modern Hopfield network in which both the questions (queries) and the memories (keys) come from the input. This last post of the series asks what happens when one side is learned instead. Two useful layers fall out: a pooling layer that summarises a whole set with a learned question, and a lookup layer with a learned memory. The pooling layer is then put to work on a classic hard problem, finding a single relevant item among many, and it turns out to need some help to get started.

1. Three Modules, One Operation

The hopfield-layers package of Ramsauer et al. (2021) offers three modules. All three compute $softmax (β Q K^{⊤}) V$ and differ only in where the queries and memories come from:

Hopfield: queries and memories both come from the input. This is attention (self- or cross-attention), as checked in the previous post.
HopfieldPooling: the query is a learned vector, a prototype of what to look for, and the memories are the elements of the input set. Whatever the size of the set, the output is one vector (or a few, with several learned queries). It is a pooling layer whose weights depend on the content of each element.
HopfieldLayer: the queries come from the input and the memories are a learned table. Each input is replaced by a mixture of learned prototypes: a lookup layer, related to key-value memory layers.

In code, HopfieldPooling(input_size=64) turns a batch of sets of shape (batch, N, 64) into (batch, 64), for any N.

2. Multiple-Instance Learning: Find the Needle

Many real problems label sets, not elements. A pathology slide is labelled “tumour” without marking the tumour cells; a patient’s immune repertoire, hundreds of thousands of receptor sequences, is labelled “has had the disease” without saying which few sequences respond to it. This is multiple-instance learning (MIL): a bag is positive if at least one instance in it is positive, and only the bag label is known.

A small, controlled version on MNIST: a bag contains $N$ handwritten digits; it is positive if it contains exactly one 9, negative if it contains none. Every model has the same convolutional encoder for single digits and a linear classifier at the end; only the way the bag’s $N$ digit embeddings $h_{1}, \dots, h_{N}$ are pooled into one vector differs:

Pooling	The bag vector	Weights
mean	$\frac{1}{N} \sum_{i} h_{i}$	equal, $1/ N$ each
max	the elementwise maximum over $i$	none: each feature takes its largest value
attention-MIL (Ilse et al., 2018)	$\sum_{i} a_{i} h_{i}$ , with $a = softmax (w^{⊤} tanh (V h_{i}))$	learned, from each element
`HopfieldPooling`	$\sum_{i} p_{i} (W_{V} h_{i})$ , with $p = softmax (β q^{⊤} W_{K} h_{i})$	learned query $q$

Attention-MIL and HopfieldPooling are close relatives: both compute a softmax over the elements from a learned scoring of each one. The Set Transformer’s “pooling by multi-head attention” (Lee et al., 2019) is the same idea again.

Training uses 500 random bags per epoch (half positive) for 30 epochs, Adam, and is scored on 500 fixed test bags built from the MNIST test set; 3 seeds per setting. The script is mil_mnist.py.

3. From Scratch: Everything Breaks as Bags Grow

Pooling	10 digits per bag	50 digits	100 digits
mean	88.8%	62.7%	50.0%
max	97.5%	94.7%	50.0%
attention-MIL	96.7%	63.3%	50.0%
`HopfieldPooling`	98.0%	50.0%	50.0%

Test accuracy, mean of 3 seeds; 50% is chance. With 10 digits, the learned poolings are the best, and HopfieldPooling puts 97% of its weight on the 9. With 50 digits, HopfieldPooling stays at chance in all three seeds, and attention-MIL ends at 50.0% and 57.2% in two of its three. With 100 digits, every method stays at chance, max pooling included.

The training loss tells why: for the failing runs it sits at $ln 2 = 0.693$ , the value of a coin flip, from the first epoch to the last. These are not overfitting or a bad learning rate; the models never start. A learned query begins near zero, so at the start every digit gets the same weight $1/ N$ : HopfieldPooling begins as mean pooling. The one 9 then contributes only $1/ N$ of the bag vector, the classifier cannot see it, and the gradient that would teach the query where to look is $N$ times smaller than in a bag of one. The query needs a classifier that already recognises 9s, and the classifier needs a query that already finds them. Max pooling avoids the trap for longer because the largest activation passes its full gradient, but at 100 digits it fails too.

Here are the bags themselves, with each digit shaded by the weight it gets. Switch between 10, 50 and 100 digits: with 10, the learned query has found the 9; with 50 and 100, HopfieldPooling gives every digit exactly $1/ N$ , and all four models answer about 50%, a pure guess:

4. With a Curriculum, the Learned Query Finds the 9

The fix is to start where learning is easy. I took each model trained on bags of 10 and continued training it for 15 epochs on the bigger bags, the same way for every pooling:

Pooling	10 digits	50 digits	100 digits
mean	88.8%	85.5%	84.3%
max	97.5%	92.6%	95.1%
attention-MIL	96.7%	91.6%	95.0%
`HopfieldPooling`	98.0%	95.5%	94.5%

Now the learned poolings work: HopfieldPooling puts 94–96% of its weight on the single 9, and ranks it first among the 50 or 100 digits in 98–99% of the positive bags. With a working start, max, attention-MIL and HopfieldPooling end up within a few points of each other. HopfieldPooling is the most consistent across seeds at 50 digits (94.6–96.2%, against 83.2–96.0% for attention-MIL and 85.8–96.4% for max), but three seeds cannot rank them further. Mean pooling stays well behind, because one 9 in 50 or 100 is diluted however well the encoder recognises it.

The same test bags as above, now with the models that started from bags of 10. The 9 lights up even among 100 digits; switch the training back to “from scratch” to compare the very same bag:

5. What the Weights Do and Do Not Tell You

Pooling weights that land on the 9 are reassuring: in this task they point at the right digit almost every time. But they are not, in general, an explanation of the decision. The classifier sees the pooled vector, and other weightings could produce a similar vector and the same answer (Jain & Wallace, 2019, showed this for attention in text models). Here we know the ground truth, the position of the 9, and can check the weights against it. In a real MIL problem, a highlighted instance is a lead to investigate, not a proof.

Two lessons carry over to real data:

The witness rate matters more than the bag size. The task gets hard when the positive instances are a tiny fraction of the bag. In DeepRC (Widrich et al., 2020), where HopfieldPooling pooled about 300,000 immune-receptor sequences per patient, what makes the problem learnable is the signal those sequences carry, not the pooling layer alone.
A learned query needs a way in. Easy examples first (small bags, or bags with more positives), a pre-trained encoder, or a warm start are the practical answers to the cold start above.

6. Common Claims, Checked

Claim	Verdict
”`HopfieldPooling` is better than attention pooling.”	They are close relatives (softmax over learned scores of the elements). With a working start they were within a few points of each other and of max pooling; `HopfieldPooling` was the most consistent across seeds here.
”A learned query finds the relevant item in any bag.”	Not from scratch: with one 9 among 50 digits it never left chance level in three of three seeds, because it starts as mean pooling. With a curriculum it found the 9 in 98–99% of positive bags.
”Max pooling is enough for multiple-instance learning.”	It was the most robust from scratch at 50 digits, but it also failed at 100 and cannot weigh several pieces of evidence against each other.
”Mean pooling is a safe default for sets.”	Not when the signal is in one element: it dilutes the needle by 1/N. With the curriculum it was the weakest method at every bag size (84–89% against 92–98%).
”The pooling weights explain the prediction.”	They pointed at the 9 here, which we could verify. In general, attention-like weights are a hint, not an explanation.
”Hopfield layers store the training data.”	`HopfieldPooling` stores a learned query and `HopfieldLayer` a learned table of prototypes; neither keeps the training examples.

The Series in One Table

Post	The idea	The number to remember
Hopfield Networks	memories are energy minima; recall rolls downhill	0.138·d random patterns, and only 5–6 correlated shapes
Modern Hopfield Networks	a steeper energy; the softmax update	capacity close to $2^{d /2}$ : 46,341 patterns in 30 dimensions
Attention Is a Hopfield Update	attention is one retrieval step, reading out values	equal to PyTorch’s attention to 3 × 10⁻⁹
Hopfield Pooling (this post)	a learned query summarises a set	94–96% of the weight on one 9 among 50–100 digits, once training gets started

Next: Pooling a Set: Hopfield Pooling against Seven Others: average, max, log-sum-exp, generalized mean, gated attention and the Set Transformer on the same task, and which ones survive bigger bags.

References

Hopfield layers and their uses

Ramsauer, H., Schäfl, B., Lehner, J., et al. (2021). Hopfield networks is all you need. ICLR. arXiv:2008.02217
Widrich, M., Schäfl, B., Pavlović, M., et al. (2020). Modern Hopfield networks and attention for immune repertoire classification. NeurIPS. arXiv:2007.13505

Pooling sets and multiple-instance learning

Dietterich, T. G., Lathrop, R. H., & Lozano-Pérez, T. (1997). Solving the multiple instance problem with axis-parallel rectangles. Artificial Intelligence, 89(1–2), 31–71.
Ilse, M., Tomczak, J. M., & Welling, M. (2018). Attention-based deep multiple instance learning. ICML. arXiv:1802.04712
Lee, J., Lee, Y., Kim, J., Kosiorek, A., Choi, S., & Teh, Y. W. (2019). Set transformer: A framework for attention-based permutation-invariant neural networks. ICML. arXiv:1810.00825

Interpreting attention

Jain, S., & Wallace, B. C. (2019). Attention is not explanation. NAACL.

Notes

Explorer