Skip to content

← Flax step 57 of 100

Medium Primitives

Squeeze-and-Excitation Block (Flax)

Jump to the task

Use image-wide information to gate channels

A squeeze-and-excitation (SE) block computes one gate per channel from the input's spatial averages. It then uses those gates to scale the original feature map. Every position in a channel receives the same gate, but the gates can change when the input image changes.

Ordinary convolutions already learn different weights for different channels. SE adds a different mechanism: input-dependent channel scaling based on a summary of the whole feature map.

Three steps and their shapes

The task takes one unbatched array x of shape (H, W, C). Let R = int(reduction) and M = max(C // R, 1).

Step Calculation Shape
Squeeze Mean over spatial axes (0, 1) (C,)
First projection Dense(M), then ReLU (M,)
Second projection Dense(C), then sigmoid (C,)
Scale Multiply the original x by the gates (H, W, C)

The two Dense layers let a gate depend on multiple channels' means. The smaller hidden width reduces the parameter count, but does not guarantee that the block avoids overfitting. Including both Dense biases, the parameter count is 2*C*M + M + C.

Integer division matters. With C=5 and R=2, the hidden width is two. With C=4 and R=16, it is one because of the clamp. With R=1, it is C, so there is no width reduction.

Gates attenuate magnitudes

Mathematically, sigmoid maps finite logits to values strictly between zero and one. Floating-point calculations can round saturated gates to zero or one. These gates do not increase an element's absolute magnitude or reverse its sign. A channel can become more prominent relative to another channel because it is attenuated less.

Each gate is independent of a sum-to-one constraint. Use sigmoid, not softmax: several channels can receive gates close to one at once. Multiplying by a gate also adds no new offset, although both Dense layers that produce the gates have biases.

A numerical example

Suppose two spatial positions contain [1, 4] and [3, 8]. Their channel means are [2, 6]. If the excitation network produces logits [0, log(3)], the gates are [0.5, 0.75]:

import jax
import jax.numpy as jnp
import numpy as np

x = jnp.array([[[1., 4.], [3., 8.]]])  # (1, 2, 2), without a batch axis
summary = x.mean(axis=(0, 1))
gates = jax.nn.sigmoid(jnp.array([0., jnp.log(3.)]))
out = x * gates
np.testing.assert_allclose(summary, [2., 6.])
np.testing.assert_allclose(gates, [0.5, 0.75])
np.testing.assert_allclose(out, [[[0.5, 3.], [1.5, 6.]]])

The chosen logits illustrate scaling only; the graded function must compute its gates through the initialized Dense layers.

Averaging removes spatial order from the gate calculation. Rearranging positions leaves the channel summaries unchanged, while the output feature map follows that same rearrangement. The gates do not identify separate important regions within a channel. Gradients also pass through the gate network, so the input derivative is not simply the gate value.

Problem

Implement se_block_forward(seed, x, reduction):

  1. Define a compact Flax SEBlock(nn.Module) with an integer reduction field. Read C = x.shape[-1].
  2. Compute jnp.mean(x, axis=(0, 1)) and hidden width max(C // reduction, 1).
  3. Apply nn.Dense(hidden_width), nn.relu, nn.Dense(C), and jax.nn.sigmoid, in that order.
  4. Multiply the original x by the resulting (C,) gates, using broadcasting over height and width.
  5. Initialize the module directly on unbatched x with jax.random.PRNGKey(int(seed)), apply it with those variables, and return the result flattened in row-major order.

Keep both Dense layers' default initializers and biases. Leave them unnamed so they become Dense_0 and Dense_1. Do not split the initialization key, add a batch axis, or wrap the block in another module.

Inputs: x is a float array (H, W, C) with positive dimensions. reduction is a positive integer-valued scalar, and seed is an integer-valued scalar. Scalars may arrive as floats; convert them to Python integers. Output: a one-dimensional array of length H * W * C.

This block has no BatchNorm or mutable state. Keep seed and reduction static when JIT-compiling the wrapper.