Squeeze-and-Excitation — the "main × gate" pattern #

The most interesting layer in this stack, structurally. It's the first op where the output is the product of two functions of the same input:

y = x ⊙ gate(x)

The "main path" is the input itself. The "gate" is a small subnetwork that computes a per-channel scaling factor by squeezing spatial info through a bottleneck. Each channel's output is multiplied by its own gate value, so SE acts as learned channel attention.

Why this matters for backprop #

This is where we hit the product rule for the first time. Both factors of x ⊙ gate(x) depend on x, so the gradient at x has contributions from both:

Through the main path: the gate itself acts as a "stop-gradient multiplier" — the gradient flowing back through the main path is just gate(x) ⊙ dy.
Through the gate path: the input flows back through the entire gate sub-network, and the cotangent it sees is x ⊙ dy (not just dy, because the gate is multiplying the main path).

The two contributions add at x. This is the same fan-in pattern as residual blocks (Residual.lean), but now driven by multiplication rather than addition — and that changes which cotangents each path sees.

Foreshadowing #

The exact same structure shows up in Transformer attention:

out = softmax(QKᵀ/√d) · V
    = attention_weights ⊙ V_with_some_extra_steps

The output is a product of "attention weights" (a function of Q and K) and V. Backprop through attention is just the product rule applied twice (once for each factor) plus the chain rule through softmax. SE is the simplest non-trivial instance of this pattern; if you understand it, attention is downhill.

What this file provides #

All foundational definitions and proofs live in Tensor.lean:

elemwiseProduct f g — pointwise product of two vector functions
elemwiseProductHasVJP — the bi-cotangent VJP (proved, no sorry)
identityHasVJP — identity backward is passthrough (proved)
pdiv_mul — product rule for partial derivatives (theorem)

This file specializes to the SE pattern: f = identity, g = gate. seBlockHasVJP takes the gate abstractly (it only needs HasVJP gate). The file then defines the concrete gate seGate (GAP → dense → swish → dense → sigmoid → broadcast) with its VJP seGateHasVJP, the broadcast it needs (broadcastFlatHasVJP; sigmoidHasVJP is in Architectures.Activations), and the full block seBlockFull / seBlockFullHasVJP.

References #

Hu, Shen, Sun 2018, Squeeze-and-Excitation Networks. https://arxiv.org/abs/1709.01507

source

noncomputable def Proofs.seBlock {n : ℕ} (gate : Vec n → Vec n) :

Vec n → Vec n

An SE-style block, parameterized by an arbitrary "gate" sub-network.

seBlock gate x = x ⊙ gate(x)

The gate can be anything — for the actual SE we use gate = sigmoid ∘ dense_exp ∘ swish ∘ dense_red ∘ globalAvgPool, but the VJP derivation doesn't care about the gate's internals. All we need is HasVJP gate.

Equations

Proofs.seBlock gate = Proofs.elemwiseProduct (fun (x : Proofs.Vec n) => x) gate

Instances For

source

noncomputable def Proofs.seBlockHasVJP {n : ℕ} (gate : Vec n → Vec n) (hg_diff : Differentiable ℝ gate) (hg : HasVJP gate) :

HasVJP (seBlock gate)

SE block VJP — direct application of the elemwise product formula.

With f = identity (so f(x) = x and f.back(x, dy) = dy), the general formula simplifies to:

back_SE(x, dy) = (gate(x) ⊙ dy) -- main path: id backward + gate.back(x, x ⊙ dy) -- gate path

First term: gradient flows back through the "main path" as gate(x) ⊙ dy — each channel scaled by its gate value.

Second term: gradient flows back through the "gate sub-network", which sees x ⊙ dy as its cotangent (not just dy!). Inside the gate, globalAvgPool will broadcast this back over spatial dims, dense_red and dense_exp will do their usual VJPs, etc.

The MLIR emits exactly this two-path backward: gate(x) * dy plus the gate's own backward chain with cotangent x * dy.

No sorry — this delegates to elemwiseProductHasVJP and identityHasVJP, both proved in Tensor.lean.

Equations

Proofs.seBlockHasVJP gate hg_diff hg = Proofs.elemwiseProductHasVJP (fun (x : Proofs.Vec n) => x) gate ⋯ hg_diff (Proofs.identityHasVJP n) hg

Instances For

What's actually inside `gate` #

For the MobileNetV3 SE block (MlirCodegen.emitSEBlock), the gate is:

Squeeze: Global average pool over (H, W) -> (B, C) g[c] = (1/(H*W)) sum_{h,w} x[c, h, w]
Reduce: Dense (C -> C/4) (or similar bottleneck) r = W_red * g + b_red
Activation: Swish r ⊙ sigma(r) (or ReLU in V3)
Expand: Dense (C/4 -> C) back to per-channel e = W_exp * sigma_swish(r) + b_exp
Sigmoid gate (or h-sigmoid in V3): sigma(e) — squashes each channel's "importance score" into [0, 1]
Broadcast back to (C, H, W) so it can multiply the main path

So the gate is actually a Vec-shaped function that takes the spatial input, summarizes it via GAP, runs it through a tiny FC network, and broadcasts the per-channel result back to spatial.

seGate at the end of this file builds gate as the composition

gate = broadcast ∘ sigmoid ∘ dense_exp ∘ swish ∘ dense_red ∘ globalAvgPool

and seGateHasVJP assembles its VJP with vjpComp from globalAvgPoolFlatHasVJP (CNN.lean), denseHasVJP, swishHasVJP and sigmoidHasVJP (Activations.lean), and broadcastFlatHasVJP (this file; broadcast is the adjoint of GAP up to the 1/(h·w) factor).

Why this generalizes #

Replace "gate" with "attention weights" and SE becomes the core of self-attention:

out = (sequence) ⊙ (per-token attention weights)

The structural pattern is identical: a main tensor multiplied by a side-computed scalar (or vector) per element. The bi-cotangent rule (elemwiseProductHasVJP) is the right tool for both. SE is the on-ramp; once you've internalized this VJP shape, attention falls out of the same theorem.

source

theorem Proofs.seBlockHasVJP_correct {n : ℕ} (gate : Vec n → Vec n) (hg_diff : Differentiable ℝ gate) (hg : HasVJP gate) (x dy : Vec n) (i : Fin n) :

(seBlockHasVJP gate hg_diff hg).backward x dy i = ∑ j : Fin n, pdiv (seBlock gate) x i j * dy j

Public correctness theorem for seBlockHasVJP: the SE-block backward (input × gate Jacobian via the product rule) equals the pdiv-contracted Jacobian of seBlock gate.

source

noncomputable def Proofs.broadcastFlat (c h w : ℕ) :

Vec c → Vec (c * h * w)

Broadcast a per-channel vector back to spatial layout. broadcastFlat c h w v idx = v (flatChannel c h w idx) — every spatial cell of channel k receives v k. This is the reindex map along flatChannel, i.e. the adjoint of globalAvgPoolFlat (up to the 1/(h·w) scale). Vec c → Vec (c*h*w).

Equations

Proofs.broadcastFlat c h w v idx = v (Proofs.flatChannel c h w idx)

Instances For

source

theorem Proofs.broadcastFlat_differentiable (c h w : ℕ) :

Differentiable ℝ (broadcastFlat c h w)

source

noncomputable def Proofs.broadcastFlatHasVJP (c h w : ℕ) :

HasVJP (broadcastFlat c h w)

Broadcast VJP — linear reindex; backward sums each channel's spatial cotangents (the adjoint of broadcast = sum-over-spatial).

Equations

One or more equations did not get rendered due to their size.

Instances For

source

noncomputable def Proofs.seGate {c h w r : ℕ} (W₁ : Mat c r) (b₁ : Vec r) (W₂ : Mat r c) (b₂ : Vec c) :

Vec (c * h * w) → Vec (c * h * w)

The squeeze-excite gate. Maps Vec (c*h*w) → Vec (c*h*w): broadcast ∘ sigmoid ∘ dense(W₂,b₂) ∘ swish ∘ dense(W₁,b₁) ∘ GAP Squeeze (GAP c·h·w → c), reduce (dense c → r), swish, expand (dense r → c), sigmoid gate, broadcast back to spatial. Every stage is smooth everywhere (swish/sigmoid smooth, dense/GAP/broadcast linear-affine), so the gate is differentiable everywhere and has a global HasVJP.

Equations