Introduction

Standard max pooling keeps only one value per pooling window. If you want the strongest k responses instead of only the single maximum, you need a custom operation because TensorFlow does not provide a built-in TopKPooling2D layer.

The right implementation depends on what you want the output to look like. Some models keep all top-k values, while others reduce them back to one value per window by averaging or summing.

What Top-k Pooling Actually Means

Suppose you have a 2x2 pooling window with values [1, 9, 4, 7]:

• max pooling keeps only 9
• top-2 pooling keeps 9 and 7
• a reduced top-2 pooling layer might return the mean, which is 8

That lets you preserve more signal than max pooling while still shrinking the spatial dimensions.

Build It with extract_patches and top_k

The general pattern in TensorFlow is:

1. extract every pooling window as a patch
2. reshape so each patch becomes a short vector
3. apply tf.math.top_k
4. optionally reduce the selected values

Here is a reusable Keras layer:

1import tensorflow as tf
2
3class TopKPooling2D(tf.keras.layers.Layer):
4    def __init__(self, pool_size=(2, 2), strides=None, k=2, reduce_mode="mean", **kwargs):
5        super().__init__(**kwargs)
6        self.pool_size = pool_size
7        self.strides = strides or pool_size
8        self.k = k
9        self.reduce_mode = reduce_mode
10
11    def call(self, inputs):
12        k_h, k_w = self.pool_size
13        s_h, s_w = self.strides
14
15        patches = tf.image.extract_patches(
16            images=inputs,
17            sizes=[1, k_h, k_w, 1],
18            strides=[1, s_h, s_w, 1],
19            rates=[1, 1, 1, 1],
20            padding="VALID",
21        )
22
23        batch = tf.shape(inputs)[0]
24        out_h = tf.shape(patches)[1]
25        out_w = tf.shape(patches)[2]
26        channels = tf.shape(inputs)[-1]
27        patch_area = k_h * k_w
28
29        patches = tf.reshape(patches, [batch, out_h, out_w, channels, patch_area])
30        values, _ = tf.math.top_k(patches, k=self.k, sorted=True)
31
32        if self.reduce_mode == "mean":
33            return tf.reduce_mean(values, axis=-1)
34        if self.reduce_mode == "max":
35            return values[..., 0]
36        if self.reduce_mode == "sum":
37            return tf.reduce_sum(values, axis=-1)
38
39        return values

Example Usage

The following example keeps the top two values from each 2x2 window and averages them:

1x = tf.constant(
2    [[[[1.0], [9.0], [2.0], [6.0]],
3      [[4.0], [7.0], [3.0], [8.0]],
4      [[5.0], [1.0], [0.0], [2.0]],
5      [[6.0], [4.0], [3.0], [9.0]]]]
6)
7
8layer = TopKPooling2D(pool_size=(2, 2), strides=(2, 2), k=2, reduce_mode="mean")
9print(layer(x).numpy())

This returns one pooled value per window, but it is based on the top two responses instead of only the maximum.

When tf.nn.top_k Alone Is Not Enough

tf.nn.top_k works well when you already have a flat vector or when you want the top k values globally. Pooling is different because you need a separate top-k operation inside every local spatial window.

That is why patch extraction is necessary. Pooling is fundamentally a local neighborhood operation, not a whole-tensor ranking operation.

Choose the Output Shape Deliberately

There are two common designs:

• reduce the selected k values to one scalar per window
• keep the k values and let the next layer consume an extra dimension

The first design is easier to integrate into ordinary CNNs. The second preserves more information, but it changes tensor shapes deeper in the model.

Common Pitfalls

• Assuming TensorFlow already has a direct top-k pooling layer like max pooling.
• Forgetting that local pooling needs per-window top-k, not global tf.nn.top_k.
• Returning all k values without updating downstream layer expectations.
• Choosing k larger than the pooling window area.
• Using a custom pooling layer without checking whether the extra cost is worth it for the task.

Summary

• Top-k pooling keeps the strongest k values from each pooling window instead of only one maximum.
• In TensorFlow, a common implementation uses tf.image.extract_patches plus tf.math.top_k.
• You can either keep all selected values or reduce them with mean, sum, or another rule.
• This is useful when max pooling discards too much local information.
• Be explicit about the output shape so later layers still receive what they expect.