Refining activation downsampling with SoftPool

In plain words

Neural image networks shrink their feature maps with pooling. The usual choices are max pooling, which keeps only the strongest value in each small patch and throws the rest away, and average pooling, which blends the patch so strong values get washed out. Both lose information that the next layers could use.

SoftPool takes a middle path. Inside each patch it gives every value a weight that grows with its strength, using an exponential followed by a normalisation so the weights add up to one. The output is the weighted sum. Strong values dominate, but weaker ones still contribute, so fine detail is less likely to vanish.

The method has no parameters to learn, and it is differentiable, so training works as usual. The authors report that swapping it in for the existing pooling layers improves classification accuracy, and that it keeps more detail when a picture is shrunk.

TL;DR

• SoftPool is a drop-in replacement for max and average pooling in CNNs, for 2D images and 3D video.
• Each output is a softmax-weighted average of the activations in its window, with the weights taken from the exponentials of those activations eq. 1 eq. 2.
• It has no trainable parameters, so parameter counts and GFLOPs stay the same as the pooling it replaces §4.4.
• Author-reported gains: best SSIM and PSNR in the downsampling tests tab. 1; ImageNet1K top-1 gains up to +1.51 (ResNet18, from scratch) across 11 architectures, with some pre-trained-swap rows not improving (see Paper vs code) tab. 2; higher HACS and Kinetics-700 top-1 for the video models tab. 7.
• Check the paper-vs-code section before implementing: the official backward pass omits a term of the exact derivative, the CPU fallback computes weights differently, and the 3D backward pass is inconsistent with its forward pass.

Problem

CNNs use pooling to shrink activation maps. This cuts computation and memory for the layers that follow, and it enlarges their receptive fields §1. The stated requirement is to lose as little information as possible while keeping the overhead low §1 §3.

Max pooling keeps one value per window and is not smooth §3.2. Average pooling weights every value equally, so a strong response can be diluted by weak neighbours §3.3. Many learned or stochastic alternatives add parameters, randomness, or extra convolutions §2 §4.4. The paper argues that this combination of lossy selection, dilution and overhead is what SoftPool targets §3.3 fig. 5.

Method

SoftPool is applied to each channel independently over a window R of k × k activations, so |R| = k² §3. The paper omits the channel index in its notation §3. Each activation a_i in R gets the softmax weight w_i, which is its exponential divided by the sum of the exponentials over R eq. 1 §3.1. The output is the weighted sum of the activations eq. 2 §3.1.

Compared with max pooling, the output is a smooth function of every activation in the window; compared with average pooling, larger activations count for more §3.1 §3.3. The window weights are shared with the backward pass: the gradient reaching each input is proportional to its forward weight §3.2 fig. 4.

The paper extends the same formulation to 3D inputs (channels, height, width, time), with the kernel becoming a spatio-temporal region §3.4.

Forward pass, one window of one channel eq. 1 eq. 2:

# SoftPool forward over one window R  [§3.1]
# Input:  window (k*k activations, one channel)
# Output: pooled scalar

def softpool_window(window):
    # Exponentiate every activation in the window  [eq. 1]
    exp_values = exp(window)
    # Normalise so the weights in the window sum to one  [eq. 1]
    weights = exp_values / sum(exp_values)
    # Output is the weighted sum of the activations  [eq. 2]
    return sum(weights * window)

Backward pass as the paper describes it §3.2 fig. 4:

# Gradient of one window, as stated in the paper  [§3.2]
# Input:  grad_out (scalar, gradient of the loss w.r.t. the output)
#         weights (k*k softmax weights saved in the forward pass)
# Output: grad_in (k*k gradients, one per activation)

def softpool_backward(grad_out, weights):
    # Each activation receives the output gradient scaled by its weight  [§3.2]
    grad_in = grad_out * weights
    return grad_in

Derived here, not stated in the paper: the exact derivative of eq. 2 with respect to each input also depends on how the weights move when that input changes. Differentiating eq. 1 and eq. 2 gives, for activation k with window output y, dy/da_k = w_k (1 + a_k − y). The backward pass above drops the (a_k − y) term. The official CUDA code also uses the weight-only form (see Paper vs code).

Numerical handling is part of the method. The authors compute in finite ranges per precision (half, single, double), clamp overflowing or underflowing intermediate values, and assign a lower limit to the weight sum so it is never zero §3.2 app. §7. These checks assume non-negative activations, as produced by ReLU app. §7.

Downsampling output size is not discussed in the paper. The official CUDA code uses no padding and drops windows that run past the edge (see Pitfalls).

Hyperparameters

Results

Author-reported. Setup: ImageNet1K for classification, with the same 11 networks trained from scratch and also pooling-replaced pre-trained networks §4.4. The downsampling tests use DIV2K, Urban100, Manga109 and Flicker2K §4.1 §4.2. Hardware for latency is listed in app. §7.

Downsampling similarity tab. 1 tab. 3. tab. 1 reports SSIM and PSNR on DIV2K, Urban100 and Manga109 for k = 2, 3, 5. SoftPool has the highest value of both metrics in all nine settings, against average, max, sum, power-average, Lp, gate, stochastic, S3, Lite-S3DPP and LIP. tab. 3 gives the Flicker2K values with latency:

Source: tab. 3. Latency is averaged over ImageNet1K; F = forward, B = backward tab. 3.

ImageNet1K classification, layer replacement tab. 2. Top-1 and top-5 in %, with gains in brackets. "Pre-trained" means pooling was swapped in a trained network without further training. "From scratch" means training with SoftPool.

Source: tab. 2. Bold is the source's bolding of the from-scratch columns. The ResNeXt50 top-5 entry (93.37) is below its original (93.70) yet is printed bold with +0.67; see Paper vs code.

Pooling-method comparison, from scratch, top-1 tab. 5:

Source: tab. 5. The Lp and Gate citation labels in tab. 5 are the reverse of those in tab. 1 and tab. 3; see Paper vs code.

Progressive replacement in InceptionV3, top-1 / top-5 tab. 6. N is the original max-pool network. Each column adds one more replaced layer: N 77.45 / 93.56; I 77.93 / 93.61; II 78.14 / 93.68; III 78.37 / 93.74; IV 78.42 / 93.78; V 78.65 / 93.84; VI 78.83 / 93.90; VII 79.04 / 93.98 (all layers replaced). Text: average single-layer gain 0.23 top-1; full-network gain +1.59 top-1 §4.5.

Action recognition tab. 7. Top-1 / top-5 in %. Models are trained on HACS, then fine-tuned for Kinetics-700 and UCF-101, except ir-CSN-101 and SlowFast (see Paper vs code). Footnote markers are the paper's.

Footnote markers tab. 7, paraphrased: (**) the authors re-implemented the model and trained it themselves; (††) the model and weights come from the authors' own repositories; (‡*) an unofficial model trained by the authors; (‡†) an unofficial repository's model with official weights; (†*) an official model trained by the authors. N/A: no model provided.

Other reported results: McNemar's test gives p well below 0.01 for the from-scratch comparisons app. §6. Parameters and GFLOPs are unchanged by the swap §4.4. The paper calls the ImageNet gains modest but important §4.

Paper vs code

Repository: alexandrosstergiou/SoftPool, owned by the first author, MIT licence. Checked at commit e11dee7e96ecad895cf871c8cbf220f7908462ed. The paper links it by a git.io short link § Abstract; the link resolves to this repository. Search: GitHub search was not needed because the link resolves.

Discrepancies, each with both citations:

1. Gradient. Paper: gradients are proportional to the forward weights §3.2. Code does exactly that: grad_in = grad_out * weight softpool_cuda_kernel.cu. It omits the (a_k − y) term of the exact derivative (see Method). The paper does not say whether the reported results used this form.
2. CPU fallback is a different operator. The non-CUDA path computes weights by summing exponentials over the channel dimension (dim=1), not within each channel's window idea.py idea.py idea.py. The CUDA path normalises per channel within the window, as eq. 1 describes. Outputs differ between devices. The fallback uses average pooling of products; the sum(kernel) factors cancel in the ratio.
3. 3D kernels. In the 3D forward pass, windows are centred (offset by kernel/2) softpool_cuda_kernel.cu. In the 3D backward pass, windows start at pd·stride with no centring softpool_cuda_kernel.cu. The gradient therefore does not match the forward window. The 3D backward also computes pooled_height = width/stride_h softpool_cuda_kernel.cu. Paper: 3D extension stated only as a kernel extension §3.4.
4. ResNet pooling window. Paper: replace the pooling layer after the first convolution §4.4. The window size is not stated for classification. Code replaces the 3 × 3, stride-2, padding-1 max pool with a 2 × 2, stride-2 SoftPool resnet.py. This changes the window size, not only the operator.
5. DenseNet pooling. Paper: replace five pooling operations, including the average pools between dense blocks §4.4. Code replaces only the stem pool densent.py. Transition layers keep AvgPool2d(2, 2); the SoftPool line is commented out densent.py.
6. Inception. Paper: replace max-pool operations within the Inception blocks §4.5. Code also replaces the average-pool branches inside each block, with a 3 × 3, stride-1 SoftPool inception.py inception.py. Paper does not mention this.
7. Training recipe. Paper: 100 epochs, LR decay every 40 epochs, batch 256, 294 → 224 crop §4.1. Code: default 90 epochs, decay ×0.1 every 30 epochs, default batch 128 train.py train.py train.py. The code also trains at 299 px train.py. Whether the reported numbers used the code defaults is not stated.
8. Video experiments. Paper: video models use the authors' implementations, and ir-CSN-101 and SlowFast use Kinetics-700 weights from other sources §4.6. The repository contains only the image models config.py; no video training code is shipped.
9. Table 2 arithmetic. Several table entries disagree with their own numbers. ResNet50 from scratch top-1: 77.35 − 76.15 = 1.20, but the table reads +1.17. ResNeXt50 from scratch top-5: 93.37 − 93.70 = −0.33, but the table reads +0.67 and bold, although it is below the original. ResNeXt101 pre-trained top-1: 78.89 − 79.31 = −0.42, but the table reads +0.58. DenseNet121 pre-trained top-1: 75.27 − 74.65 = +0.62, but the table reads +0.57 tab. 2.
10. Citation labels. In tab. 5, Lp is labelled with the Gate citation and Gate with the Lp citation; tab. 1, tab. 3 and the text assign them the other way round §2 tab. 3.

Pitfalls

• Backward is not the exact gradient. The official backward pass uses the weights alone. An implementer who wants the exact gradient of eq. 2 must add the (a_k − y) term; the two give different training dynamics. Test both against autograd before choosing (see Method).
• Overflow in exponentials. The CUDA kernels compute exp(a) without subtracting the window maximum, and clamp overflowing sums to the type's maximum softpool_cuda_kernel.cu softpool_cuda_kernel.cu. Derived: in float16 (the paper's testing precision app. §7) exp(a) exceeds the largest finite half value for a above about 11.1. Activations that large saturate. A max-subtracted formulation gives the same weights but is not what the code does.
• Sign handling. The CUDA clamp keeps the sign of the output (it clamps magnitude) softpool_cuda_kernel.cu. The overflow and underflow checks assume non-negative (ReLU) activations app. §7. Negative inputs are not covered by the paper's checks.
• Window handling without padding. The CUDA kernel uses no padding and drops windows that run past the edge. The output shape is computed in Python as floor((H − k)/s) + 1 idea.py. The kernel also computes the pooled height as H/s, which can differ from this shape. The output buffer is zero-initialised, so unwritten cells stay zero. Verify shapes if the input size is not a multiple of the stride.
• Stride is unstated for the downsampling tests. Use the code default (stride equal to kernel) only as a stated assumption idea.py.
• Small margins in the similarity tests. SoftPool wins on every metric, but some margins are small: for example SSIM 0.729 against 0.714 for average pooling on DIV2K at k = 2 tab. 1. These are image-quality metrics, not classification accuracy.
• Single runs and significance. The number of seeds per setting is not stated. McNemar's test compares paired predictions of two trained models app. §6; it does not measure seed variance. Gains of about 0.5 to 1.0 top-1 may be within seed-to-seed spread for these architectures. Plan multi-seed runs before drawing conclusions.
• Unfair video comparison. The baselines are pre-trained on different data (ir-CSN-101 on IG65M, SlowFast on ImageNet) and the SoftPool models use the same pre-training as their baselines §4.6. Comparisons across rows are not controlled.
• Latency. Latencies come from one 2080 Ti setup in float16 app. §7. The "Latency" column is averaged over ImageNet1K while the similarity columns use Flickr2K tab. 3. Do not transfer the numbers to other hardware.
• Parameter claim. Parameters and FLOPs are unchanged, but the official DenseNet and ResNet replacements change the window size (see Paper vs code, items 4 and 5). A parameter-matched comparison still needs the same window size.

Lineage

• Builds on: the max-like cortex model of Riesenhuber and Poggio and the pooling study of Boureau et al., both cited for the softmax-like design §3.1. Earlier pooling operators compared: max and average §2, stochastic [zeiler2013stochastic], S3 [zhai2017s3pool], Lp [gulcehre2014learned], gate [lee2016generalizing], DPP [saeedan2018detail], LIP [gao2019lip]. arXiv ids for these are not listed in this card and they are not linked here.
• Compared against: all of the above in tab. 1, tab. 3 and tab. 5.
• Follow-up: the same authors' AdaPool work is named in the repository README README.md. Not compiled here.