Pooling a Set: Hopfield Pooling against Seven Others

Hopfield Pooling compared four ways to pool a bag of digit embeddings into one vector. This post puts HopfieldPooling against seven alternatives, from plain average pooling to the Set Transformer, and asks two questions the first comparison left open:

Train small, test big. A pooling trained on bags of 10 digits is applied, unchanged, to bags of up to 200. Which ones still find the 9?
More evidence. What changes when a positive bag holds three or ten 9s instead of one?

Same task as before (a bag is positive if it contains a 9), same encoder, same training, three seeds each. Only the pooling differs.

Try it yourself All three experiments are in one self-contained notebook: it downloads MNIST, installs hopfield-layers, defines the eight poolings and reproduces the tables below. In Colab, choose Runtime → Change runtime type → GPU, then Runtime → Run all. On a CPU with two threads it gives exactly the numbers of this post; on a GPU they can differ slightly. The file is also in the repository: pooling_comparison.ipynb.

1. Eight Ways to Pool a Set

Pooling	Bag vector from element embeddings $h_{1}, \dots, h_{N}$	Parameters (pooling only)
mean (average pooling)	$\frac{1}{N} \sum_{i} h_{i}$	0
max	elementwise $max_{i} h_{i}$	0
log-sum-exp	$\frac{1}{r} lo g \frac{1}{N} \sum_{i} e^{r h_{i}}$ , learned temperature $r$	1
generalized mean (GeM)	$(\frac{1}{N} \sum_{i} h_{i}^{p})^{1/ p}$ , learned exponent $p$	1
attention-MIL (Ilse et al., 2018)	$\sum_{i} a_{i} h_{i}$ , $a = softmax (w^{⊤} tanh (V h_{i}))$	2,113
gated attention (Ilse et al., 2018)	as above with $tanh (V h_{i}) ⊙ σ (U h_{i})$	4,193
Set Transformer PMA (Lee et al., 2019)	a learned seed attends to the set with 4 heads, then LayerNorm and an MLP	25,280
`HopfieldPooling` (Ramsauer et al., 2021)	a learned query attends to the set (1 head), with LayerNorm on the inputs	8,800

The code

Each pooling takes the digit embeddings h of shape (bags, N, d) and returns one vector per bag, and the per-digit weights when it computes any. This is the code of the notebook (where the Hopfield class also has a switch to turn off its input LayerNorms, used in section 4):

class Mean(nn.Module):
    def forward(self, h): return h.mean(1), None
 
class Max(nn.Module):
    def forward(self, h): return h.max(1).values, None
 
class LogSumExp(nn.Module):          # (1/r) log mean exp(r h): mean as r -> 0, max as r -> infinity
    def __init__(self): super().__init__(); self.log_r = nn.Parameter(torch.zeros(1))
    def forward(self, h): r = self.log_r.exp(); return (torch.logsumexp(r * h, 1) - math.log(h.shape[1])) / r, None
 
class GeM(nn.Module):                # (mean h^p)^(1/p): mean at p = 1, max as p -> infinity
    def __init__(self): super().__init__(); self.p = nn.Parameter(torch.tensor(3.0))
    def forward(self, h): p = self.p.clamp(min=1.0); return h.clamp(min=1e-6).pow(p).mean(1).pow(1.0 / p), None
 
class AttentionMIL(nn.Module):       # Ilse et al. (2018): a = softmax(w^T tanh(V h)), optionally gated by sigmoid(U h)
    def __init__(self, gated=False):
        super().__init__(); self.V, self.w = nn.Linear(D, 32), nn.Linear(32, 1); self.U = nn.Linear(D, 32) if gated else None
    def forward(self, h):
        s = torch.tanh(self.V(h))
        if self.U is not None: s = s * torch.sigmoid(self.U(h))
        a = torch.softmax(self.w(s).squeeze(-1), 1); return (a.unsqueeze(-1) * h).sum(1), a
 
class PMA(nn.Module):                # Set Transformer: H = LN(S + MHA(S, X, X)), Z = LN(H + rFF(H)), one learned seed S
    def __init__(self, heads=4):
        super().__init__(); self.S = nn.Parameter(torch.randn(1, 1, D) * 0.02); self.mha = nn.MultiheadAttention(D, heads, batch_first=True)
        self.ln1, self.ln2 = nn.LayerNorm(D), nn.LayerNorm(D); self.ff = nn.Sequential(nn.Linear(D, D), nn.ReLU(), nn.Linear(D, D))
    def forward(self, h):
        S = self.S.expand(h.shape[0], 1, -1); att, a = self.mha(S, h, h, need_weights=True, average_attn_weights=True)
        H = self.ln1(S + att); return self.ln2(H + self.ff(H)).squeeze(1), a[:, 0]
 
class Hopfield(nn.Module):           # HopfieldPooling (Ramsauer et al., 2021): a learned query, softmax(beta q K^T) V
    def __init__(self):
        super().__init__(); self.hp = HopfieldPooling(input_size=D, hidden_size=32, quantity=1)
    def forward(self, h):
        return self.hp(h), (None if self.training else self.hp.get_association_matrix(h)[:, 0, 0])

Everything else is shared: a small CNN maps each 28 × 28 digit to a 64-dimensional embedding $h_{i}$ , the pooling turns the bag into one vector, and one linear layer gives the logit for “contains a 9”. Training is Adam (learning rate $5 \cdot 1 0^{- 4}$ , weight decay $1 0^{- 4}$ ) on 500 random bags per epoch for 30 epochs, with the binary cross-entropy loss.

Log-sum-exp and the generalized mean are dials: at $r \to 0$ or $p = 1$ they are the mean, and as $r$ or $p$ grows they approach the max. The last four compute a weight for every element, so they can say which digit they looked at.

2. On Bags of 10, Everything Works

Pooling	Accuracy, bags of 10
mean	88.8%
max	97.5%
log-sum-exp	97.6%
generalized mean	97.1%
attention-MIL	96.7%
gated attention	97.6%
Set Transformer PMA	96.6%
`HopfieldPooling`	98.0%

Every method except mean pooling lands between 96.6% and 98.0%, within seed-to-seed variation of each other. The interesting differences appear only when the conditions change.

3. Train Small, Test Big

Each model trained on bags of 10 is now tested, without any retraining, on bags of 25, 50, 100 and 200 digits with one 9:

Pooling (trained on 10)	tested on 10	on 100	on 200 digits
mean	86.1%	50.4%	50.0%
max	95.1%	90.8%	87.1%
log-sum-exp	93.7%	93.4%	92.9%
generalized mean	94.5%	88.9%	76.7%
attention-MIL	94.9%	91.2%	93.6%
gated attention	94.9%	93.9%	93.7%
Set Transformer PMA	94.2%	54.1%	50.0%
`HopfieldPooling`	95.5%	92.9%	91.8%

(The “tested on 10” column uses a different set of test bags than section 2, hence slightly different numbers.)

Three groups emerge:

Collapse to chance: mean and the Set Transformer’s PMA. Mean pooling dilutes the 9 by $1/ N$ ; what worked with 10 digits is invisible among 200. PMA collapses too, although it is the largest pooling here. Its attention explains why: even on bags of 10 it puts only 14% of its weight on the 9, barely above the uniform 10%, so in practice it acts like mean pooling followed by a learned network.
Slow decline: max and the generalized mean. Max loses 8 points by 200 digits, as more non-9 digits produce large activations by chance. The generalized mean, whose exponent stayed at $p \approx 3.1$ (it started at 3), is closer to the mean and declines faster.
Holding up: log-sum-exp, attention-MIL, gated attention and HopfieldPooling, all between 91.8% and 93.7% at 200 digits. The softmax-weighted poolings still find the 9 because its score stands out: HopfieldPooling’s weight on the 9 falls from 95% at 10 digits to 69% at 200, as the softmax spreads over more elements, but stays far above the $1/200 = 0.5%$ of a uniform weighting. Log-sum-exp, a single learned number ( $r \approx 1.1$ ), generalises about as well as the learned scoring networks.

4. More Evidence: One, Three or Ten 9s

The previous post found that, trained from scratch, most poolings never start learning on bags of 50 with one 9. Is that about the bag size, or about how rare the evidence is? Here the bags have 50 digits, and positive bags contain 1, 3 or 10 nines:

Pooling (bags of 50, from scratch)	one 9 (2%)	three 9s (6%)	ten 9s (20%)
mean	62.7%	99.4%	100%
max	94.7%	99.5%	100%
log-sum-exp	64.9%	99.6%	100%
generalized mean	94.3%	99.6%	99.9%
attention-MIL	63.3%	99.1%	100%
gated attention	76.8%	99.8%	100%
Set Transformer PMA	50.0%	50.0%	100%
`HopfieldPooling`	50.0%	50.0%	100%

With three 9s, even average pooling reaches 99.4%: 6% of the bag is enough signal for a mean to detect. The witness rate, not the bag size, decides whether training can start. Two poolings are the exception: PMA and HopfieldPooling stay at chance with three 9s in all three seeds, and only start with ten.

Why are the two learned-query poolings the slowest to start? At initialisation their query weighs every element equally, so they begin as mean pooling, which with three 9s does work. One difference is that both normalise with LayerNorm: HopfieldPooling normalises its inputs by default, and PMA its output. Normalisation removes the overall size of the activations, which is one of the things a mean uses to detect “some of these digits look like 9s”. With its input LayerNorms switched off, HopfieldPooling learned in one of three seeds (99.6%, with 73% of its weight on the 9s); the other two stayed at 50%. So normalisation explains part of the slow start, not all of it, and I did not isolate the rest.

5. What to Use

If your sets…	Then
carry the signal in a few elements, and you need to know which	a softmax-weighted pooling: attention-MIL, gated attention or `HopfieldPooling`; start training on small or evidence-rich bags
will be larger at test time than in training	avoid mean pooling and, in this test, PMA; log-sum-exp, the attention poolings and `HopfieldPooling` held up best
need a cheap, robust baseline	max pooling, or log-sum-exp with its single learned temperature
carry the signal in many elements	anything works, including average pooling

Three seeds per setting on one synthetic task are enough to see these patterns, not to rank methods within a few points of each other. On bags of 10, for example, HopfieldPooling (98.0%) and gated attention (97.6%) are not distinguishable.

6. Common Claims, Checked

Claim	Verdict
”Average pooling is not enough for multiple-instance learning.”	It depends on the witness rate: with one 9 in 50 digits it failed (62.7%), with three it reached 99.4%.
”Attention-based pooling generalises to larger sets.”	The softmax-weighted poolings did (91.8–93.7% at 200 digits after training on 10), but the Set Transformer’s PMA collapsed to chance, like mean pooling. Having attention is not enough; it has to concentrate on the relevant elements.
”`HopfieldPooling` is the best pooling.”	Not in this comparison: it was the best on bags of 10 (98.0%) and held up when tested on bigger bags, but within a few points of the attention poolings and log-sum-exp, and it was among the slowest to start learning from rare evidence.
”More parameters make a better pooling.”	PMA, with 25,280 parameters, generalised worst; log-sum-exp, with one, generalised among the best.
”Max pooling is a strong baseline.”	Yes: the most reliable from scratch with one 9 in 50 (94.7%), and it declined gracefully on bigger bags (87.1% at 200).
”Learned pooling parameters adapt to the task.”	The two learned dials barely moved: log-sum-exp’s temperature went from 1 to about 1.1, the generalized mean’s exponent from 3 to about 3.1. Their behaviour comes mostly from the starting value.

Summary

Question	Answer from these experiments
Who works on small bags?	everyone except average pooling, within seed variation
Who survives bigger test bags?	log-sum-exp, attention-MIL, gated attention, `HopfieldPooling`; max declines; mean and PMA collapse
What decides whether training starts?	the share of positive elements: 2% fails for most methods, 6% works for most, 20% for all
Where does `HopfieldPooling` stand?	top group on small bags and on bigger test bags, but slow to start from rare evidence, partly because of its input normalisation

References

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
Ramsauer, H., Schäfl, B., Lehner, J., et al. (2021). Hopfield networks is all you need. ICLR. arXiv:2008.02217
Radenović, F., Tolias, G., & Chum, O. (2019). Fine-tuning CNN image retrieval with no human annotation. IEEE TPAMI, 41(7), 1655–1668. (generalized-mean pooling)
Zaheer, M., Kottur, S., Ravanbakhsh, S., Póczos, B., Salakhutdinov, R. R., & Smola, A. J. (2017). Deep sets. NeurIPS. (permutation-invariant pooling of sets)

Notes

Explorer