Max pooling
You will learn
What max pooling keeps, where it remembers taking it from, and where the gradient goes.
Max pooling keeps the largest value of each window and drops the rest, so the network cares less about exactly where a feature sits. On the spec's 4 x 4 grid, 2 x 2 windows with stride 2 keep 4, 8, 9 and -2, taken from cells 5, 6, 13 and 11. Overlapping 3 x 3 windows with stride 1 keep 8, 8, 9 and 9. In the backward pass each window's whole gradient goes to its maximum: gradients 1, 2, 3 and 4 land on cells 5, 6, 13 and 11, and every other cell gets 0. The spec maxpool2d_layer.t27 computes the outputs, their positions and the gradient.
Try it
Switch to 3 x 3 windows and find the cell two outputs share; then turn on the gradient for 2 x 2 windows.
Pick a window size and stride and tap an output: max pooling keeps the largest value of its window and remembers where it was, and the gradient later goes to that one cell.
specs/ml/layers/maxpool2d_layer.t27
// SPDX-License-Identifier: Apache-2.0
// t27/specs/
// Reduce spatial dimensions | φ² + 1/φ² = 3 | TRINITY
// #8980: forward and backward compared two values and threw them away
// (placeholders), and every test asserted `result != undefined`. Max pooling now
// runs over one H x W channel (row-major) with a k x k window and a stride. It
// records each window's argmax, and backward sends each gradient back to that
// element only. Padding and dilation are carried in the config but not read: the
// layer assumes padding 0 and dilation 1.
module Maxpool2d;
use base::types;
use math::constants;
// ═══════════════════════════════════════════════════════════
// 2. Types
// ═══════════════════════════════════════════════════════════
pub const MaxPool2DConfig = struct {
kernel_size : u32,
stride : u32,
padding : u32,
dilation : u32,
return_indices : bool,
};
// ═══════════════════════════════════════════════════════════
// 3. Core Functions
// ═══════════════════════════════════════════════════════════
// Output side for an input side n: (n - k) / stride + 1
fn out_size(n: u32, config: MaxPool2DConfig) -> u32 {
return (n - config.kernel_size) / config.stride + 1;
}
// output[oy][ox] = max of the k x k window at (oy * stride, ox * stride);
// indices[oy][ox] = the flat input index that held it
fn forward(input: []const f32, height: u32, width: u32, config: MaxPool2DConfig, output: []f32, indices: []u32) -> void {
const oh = out_size(height, config);
const ow = out_size(width, config);
var oy : u32 = 0;
while (oy < oh) {
var ox : u32 = 0;
while (ox < ow) {
var best_index : u32 = (oy * config.stride) * width + ox * config.stride;
var ky : u32 = 0;
while (ky < config.kernel_size) {
var kx : u32 = 0;
while (kx < config.kernel_size) {
const idx : u32 = (oy * config.stride + ky) * width + ox * config.stride + kx;
if (input[idx] > input[best_index]) {
best_index = idx;
}
kx = kx + 1;
}
ky = ky + 1;
}
output[oy * ow + ox] = input[best_index];
indices[oy * ow + ox] = best_index;
ox = ox + 1;
}
oy = oy + 1;
}
}
// grad_input[indices[o]] += grad_output[o]; every other input gets 0
fn backward(grad_output: []const f32, indices: []const u32, grad_input: []f32) -> void {
var i : usize = 0;
while (i < grad_input.len) {
grad_input[i] = 0.0;
i = i + 1;
}
var o : usize = 0;
while (o < grad_output.len) {
grad_input[indices[o]] = grad_input[indices[o]] + grad_output[o];
o = o + 1;
}
}
fn cfg(kernel: u32, stride: u32) -> MaxPool2DConfig {
return MaxPool2DConfig{ .kernel_size = kernel, .stride = stride, .padding = 0, .dilation = 1, .return_indices = true };
}
// The output side on the 4 x 4 fixture for a k x k window and a stride.
// out_size takes a MaxPool2DConfig struct and does not lower to WASM; this
// takes two integers, so a host reads the side from the compiled spec.
fn out_side(kernel: u32, stride: u32) -> u32 {
return out_size(4, cfg(kernel, stride));
}
// The 4 x 4 fixture
// 1 2 5 6
// 3 4 8 7
// 0 -1 -3 -2
// -4 9 -6 -5
fn pool_at(kernel: u32, stride: u32, which: u8, index: usize) -> f32 {
const input : [16]f32 = [1.0, 2.0, 5.0, 6.0, 3.0, 4.0, 8.0, 7.0, 0.0, -1.0, -3.0, -2.0, -4.0, 9.0, -6.0, -5.0];
var output : [9]f32 = [0.0] ** 9;
var indices : [9]u32 = [0] ** 9;
forward(&input, 4, 4, cfg(kernel, stride), &output, &indices);
if (which == 0) {
return output[index];
}
return @as(f32, @floatFromInt(indices[index]));
}
// 2 x 2, stride 2, gradients (1, 2, 3, 4) for the four windows; grad_input[index]
fn backward_at(index: usize) -> f32 {
const input : [16]f32 = [1.0, 2.0, 5.0, 6.0, 3.0, 4.0, 8.0, 7.0, 0.0, -1.0, -3.0, -2.0, -4.0, 9.0, -6.0, -5.0];
var output : [4]f32 = [0.0] ** 4;
var indices : [4]u32 = [0] ** 4;
forward(&input, 4, 4, cfg(2, 2), &output, &indices);
const grad_output : [4]f32 = [1.0, 2.0, 3.0, 4.0];
var grad_input : [16]f32 = [9.0] ** 16;
backward(&grad_output, &indices, &grad_input);
return grad_input[index];
}
// ═══════════════════════════════════════════════════════════
// TDD: Tests (from .tri behaviors)
// ═══════════════════════════════════════════════════════════
// 2 x 2, stride 2: windows' maxima are 4, 8, 9, -2
test forward_basic_case
given first = pool_at(2, 2, 0, 0)
when last = pool_at(2, 2, 0, 3)
then first == 4.0 and pool_at(2, 2, 0, 1) == 8.0 and pool_at(2, 2, 0, 2) == 9.0 and last == -2.0
// The argmax indices: 4 at 5, 8 at 6, 9 at 13, -2 at 11
test forward_max_finding
given first = pool_at(2, 2, 1, 0)
when third = pool_at(2, 2, 1, 2)
then first == 5.0 and pool_at(2, 2, 1, 1) == 6.0 and third == 13.0 and pool_at(2, 2, 1, 3) == 11.0
// 3 x 3, stride 1 gives a 2 x 2 output whose windows overlap
test forward_overlapping_windows
given first = pool_at(3, 1, 0, 0)
when last = pool_at(3, 1, 0, 3)
then first == 8.0 and pool_at(3, 1, 0, 1) == 8.0 and pool_at(3, 1, 0, 2) == 9.0 and last == 9.0
test backward_basic_case
given at_max = backward_at(5)
when elsewhere = backward_at(0)
then at_max == 1.0 and elsewhere == 0.0
// Each window's gradient lands on its argmax only: the sum is conserved
test backward_gradient_propagation
given total = backward_at(5) + backward_at(6) + backward_at(13) + backward_at(11)
when others = backward_at(0) + backward_at(1) + backward_at(4) + backward_at(15)
then total == 10.0 and others == 0.0 and backward_at(13) == 3.0
// The four window/stride pairs the nn-maxpool widget offers, on the 4 x 4 fixture:
// 2 x 2 stride 2 -> 2, 3 x 3 stride 1 -> 2, 2 x 2 stride 1 -> 3, 4 x 4 stride 1 -> 1
test out_side_of_each_widget_config
given halved = out_side(2, 2)
when whole = out_side(4, 1)
then halved == 2 and out_side(3, 1) == 2 and out_side(2, 1) == 3 and whole == 1
// Edges: a 1 x 1 window at stride 1 keeps the side; a stride that does not
// divide (4 - k) rounds down (3 x 3 at stride 2 fits one window, 1 x 1 at
// stride 3 fits two)
test out_side_identity_and_floor
given same = out_side(1, 1)
when floored = out_side(3, 2)
then same == 4 and floored == 1 and out_side(1, 3) == 2
invariant maxpool_output_size_formula
assert out_size(4, cfg(2, 2)) == 2 and out_size(4, cfg(3, 1)) == 2 and out_size(5, cfg(2, 2)) == 2
All lessons
Module 1 · One neuron
A neuron sums its inputs and bends the sum: ReLU, the sigmoid and tanh, and where each one stops passing a gradient back.
Module 2 · New bends
The smooth curves of current models: GELU with its dip, Swish with its knob, and ELU with a floor instead of a cliff.
Module 3 · A layer
What sits around the neurons: softmax turns scores into a choice, an embedding turns a word into numbers, dropout switches units off.
Module 4 · Pictures
Reading a picture in patches: a window averaged over its patch, max pooling, and batch normalisation of features on different scales.
Module 5 · The error
How wrong a guess is, as one number: mean squared error, cross-entropy for probabilities, and the Huber loss for outliers.
Module 6 · Downhill
Walking down the loss: gradient descent with momentum, RMSProp that ignores the gradient's size, and Adam that combines them.
Module 7 · Learning
Where the gradients come from and what they achieve: backpropagation through a layer, a network that learns XOR, and LAMB for large batches.
Module 8 · Memory and attention
Networks that read a sequence: a recurrent cell and its fading echo, the gates of an LSTM, and attention over every position at once.
Module 9 · The transformer and the reward
The parts of a transformer and of an agent: positions as clocks, one encoder block, and PPO, which learns from reward without lurching.