7 EfficientNet

A convolution treats every channel equally

By the time a feature map reaches a deep stage of a CNN, each channel is a learned “feature detector” for something specific: a headlight, a curve, a fur texture, a leaf edge. A standard \(3 \times 3\) convolution that consumes this feature map treats every one of those channels as equally important contributors to every output channel. Its kernel learns a fixed mixing recipe that’s applied to every spatial position of every input image, regardless of what’s actually in the image.

That’s wasteful in a particular way. A “headlight” channel is useful when there’s a car in the image and useless on a beach. A “leaf-edge” channel is useful in a tree-heavy region and useless in the sky. The network has no built-in way to say “for this image, dial channel 47 up and channel 112 down.” Hu et al. 2018 (arXiv:1709.01507) designed a tiny module that lets the network do exactly that: look at the global content of the feature map, produce a per-channel gain, and rescale every channel by its gain before passing the feature map forward.

That module is the Squeeze-and-Excitation block, and at its core it is a learned per-channel attention mechanism, contemporary with the transformer’s (Vaswani et al. 2017) and three years ahead of attention’s arrival in vision.

7.1 Run it first

Before the attention mechanism, train the network that uses it. Four commands, and about forty minutes of GPU time:

lake exe cache get                            # Mathlib oleans, ~30 s
./scripts/datasets/download_imagenette.sh     # 330 MB dl, 2.4 GB unpacked
lake build efficientnet-verified-adam
./.lake/build/bin/efficientnet-verified-adam data

On one RTX 4060 Ti (CUDA 12.9), from runs/2026-08-12-enet-imagenette-xla-cuda/. XLA’s startup banner is removed, the network description line is wrapped, and the middle of the run is elided:

[pjrt_ffi] XLA backend: PJRT 0.112, 1 device(s)
[pjrt_ffi] compiled verified_mlir/efficientnet_adam_train_step.mlir
             (@efficientnet_adam_train_step, 740 outputs, 1 replica)
             in 10603 ms
[pjrt_ffi] compiled verified_mlir/efficientnet_fwd.mlir
             (@efficientnet_fwd, 1 outputs, 1 replica) in 635 ms
[pjrt_ffi] compiled verified_mlir/efficientnet_fwd_eval.mlir
             (@efficientnet_fwd_eval, 1 outputs, 1 replica) in 454 ms
EfficientNet-B0 on Imagenette 224² (stem-s2 → 16 MBConv [t,c,n,s,k],
  swish + squeeze-excite + batch-norm, 5 downsamples 224→7 →
  head 320→1280 → GAP → dense)
  via the VERIFIED renderer → XLA/PJRT → GPU
  train 9469, val 3925; bs 32, EfficientNet-B0 adam
    (cosine+warmup 3ep, baseLR 0.001000), He init
  running-stats BN: 49 layers, 42016 stat floats
    → eval via @efficientnet_fwd_eval
Epoch 1/80: loss=2.063774 lr=0.000333
  epoch 1: val_acc = 758/3925 = 19.312102%  top5 = 2630/3925 = 67.006369%
Epoch 2/80: loss=1.631765 lr=0.000667
  epoch 2: val_acc = 1216/3925 = 30.980892%  top5 = 2889/3925 = 73.605096%
Epoch 3/80: loss=1.430369 lr=0.001000
  epoch 3: val_acc = 1526/3925 = 38.878981%  top5 = 3173/3925 = 80.840764%
Epoch 4/80: loss=1.294553 lr=0.001000
  epoch 4: val_acc = 2116/3925 = 53.910828%  top5 = 3521/3925 = 89.707006%
Epoch 79/80: loss=0.504706 lr=0.000000
  epoch 79: val_acc = 3535/3925 = 90.063694%  top5 = 3857/3925 = 98.267516%
Epoch 80/80: loss=0.505049 lr=0.000000
  epoch 80: val_acc = 3534/3925 = 90.038217%  top5 = 3858/3925 = 98.292994%
done (trained EfficientNet-B0 adam + cosine/warmup via packed threading).

Eighty epochs at about 30 seconds each, forty-one minutes in total, and 90.04% top-1 with 98.29% top-5 on Imagenette’s 3,925-image validation split. The best epoch reached \(90.17\% \). That is level with ResNet-34’s five-seed \(89.99 \pm 0.32\) and eight tenths of a point over MobileNetV2’s \(89.25\% \), on the same data, the same recipe and the same eighty epochs. EfficientNet-B0 does it with \(5.3\times \) fewer parameters than ResNet-34 and in forty-one minutes against forty-seven.

Three steps, one tensor product

An SE block on a feature map \(x \in \mathbb {R}^{B \times C \times H \times W}\) runs three operations in sequence:

1. Squeeze. Global-average-pool across spatial dimensions: \(s_c = \frac{1}{HW} \sum _{i,j} x_{c, i, j}\). Result is a \(B \times C\) tensor: one scalar per image per channel, summarizing the channel’s overall activity.

2. Excite. Push \(s\) through a tiny two-layer MLP with a bottleneck: \(g = \sigma (W_2\, \mathrm{ReLU}(W_1 s))\), where \(W_1\) is \(C \to C/r\) and \(W_2\) is \(C/r \to C\) for a reduction ratio \(r\) (typically 16). The sigmoid squashes each output to \([0, 1]\). Result: \(g \in \mathbb {R}^{B \times C}\), a per-channel gate value.

3. Apply. Multiply the original feature map by the gates, broadcast across the spatial axes: \(y_{c, i, j} = g_c \cdot x_{c, i, j}\). Channels with gate near 1 pass through, and channels with gate near 0 are suppressed.

That’s it. The MLP operates on a \([B, C]\) tensor, not \([B, C, H, W]\), so it’s tiny, at \(C^2 / r + C/r + C\) parameters per block against millions for a standard conv at the same width.

From the framework: it’s an elementwise product

What is the SE block, as a function? Up to a learned scalar factor per channel, it’s just an elementwise product of the input with a side branch:

where the gate \(g\) is built from the saved input \(x\) via two dense layers and a sigmoid. We already proved the VJP of an elementwise product (Theorem 13). We already proved the VJP of dense (Theorem 19). We already proved the VJP of identity, and an additive fan-in collapses to the same product-rule structure as elementwise. Theorem 96 is one line of composition, and it introduces no new analytic dependencies.

Attention in Disguise

SE looks at the entire feature map, produces a content-dependent re-weighting, and applies it back to the data. That is a learned attention operation in the standard sense. The reduction is over the spatial axes (so each token-equivalent is a channel, not a patch), but the shape is the same: a query (GAP summary) attending over a collection (the channels) to produce weights, weights gating the original signal.

The transformer’s self-attention (Vaswani et al. 2017) generalizes the same pattern from a fixed query (the GAP) to a learned sequence of queries (the token positions) and from a fixed key/value (the channels) to a learned key/value (per-token features). SE is the bare minimum version of that idea. Most post-2018 convolutional architectures carry SE or one of its descendants (CBAM, ECA, GE), because the accuracy-per-parameter cost is low and the operation composes cleanly with anything else.

EfficientNet-B0: MobileNet-V2 with SE everywhere

EfficientNet-B0 (Tan & Le 2019, arXiv:1905.11946) is, at the architectural level, MobileNet-V2 with two changes:

1. Every MBConv block contains an SE sub-module after the depthwise conv. The verified spec spells that as its own constructor, .mbConvSENB, rather than as a flag.

2. Swish (\(x \cdot \sigma (x)\)) replaces ReLU as the activation. Smoother gradient, modest empirical lift, and no new VJP structure, just chain rule through a sigmoid.

The paper’s actual thesis is compound scaling: jointly scale depth, width, and input resolution by a single coefficient \(\phi \) and you get the EfficientNet family B0 through B7, with B7 reaching ImageNet state-of-the-art at the time. At the B0 baseline (the one in this book), the architectural lift over MobileNet-V2 is essentially SE plus Swish. The compound scaling story lives in the B0\(\to \)B7 progression, which we don’t fit on Imagenette because the dataset is too small for the return-on-scaling to show up.

7.2 The theorems

Figure 7-1 is the dependency graph of this section.

The ladder for EfficientNet-B0: the VJP, the chain, the step, the loss gradient and the data-parallel tie; no seal.

7.3 Example: EfficientNet-B0 on Imagenette

The architecture

EfficientNet-B0 is the MobileNetV2 column with Squeeze-and-Excitation in every block: the same \(3 \times 3\) stride-2 stem, seven stages of MBConv blocks (kernel sizes mixing \(3 \times 3\) and \(5 \times 5\) depthwise), the \(1 \times 1\) lift to 1280, global average pool, and dense head. Figure 7-2 is the network, and its inset shows one MBConv block, an inverted residual with an SE gate (highlighted) inserted before the project.

The verified trainer

The pattern is §5-2’s; only the layers list and the bnChannels that follows from it are new:

import LeanMlir.Verified.NetsCore
import LeanMlir.Verified.Train

-- 1. The network. `slug` names the committed render, `bnChannels`
--    drives BN threading, `NB` means the BN-followed convs carry
--    no bias.
def efficientnetVerified : VerifiedNetSpec where
  name     := "EfficientNet-B0"
  slug     := "efficientnet"
  inC      := 3
  imageH   := 224
  imageW   := 224
  data     := .imagenette
  layers   := [
    .convBnNB 3 32 3 2,             -- stem 3x3-s2   224->112
    .mbConvSENB  32   32  16  8 3,  -- stage 1, t=1 (no expand)
    .mbConvSENB  16   96  24  4 3,  -- stage 2       112->56
    .mbConvSENB  24  144  24  6 3,
    .mbConvSENB  24  144  40  6 5,  -- stage 3        56->28
    .mbConvSENB  40  240  40 10 5,
    .mbConvSENB  40  240  80 10 3,  -- stage 4        28->14
    .mbConvSENB  80  480  80 20 3,
    .mbConvSENB  80  480  80 20 3,
    .mbConvSENB  80  480 112 20 5,  -- stage 5           @14
    .mbConvSENB 112  672 112 28 5,
    .mbConvSENB 112  672 112 28 5,
    .mbConvSENB 112  672 192 28 5,  -- stage 6        14->7
    .mbConvSENB 192 1152 192 48 5,
    .mbConvSENB 192 1152 192 48 5,
    .mbConvSENB 192 1152 192 48 5,
    .mbConvSENB 192 1152 320 48 3,  -- stage 7            @7
    .convBnNB 320 1280 1 1,         -- head, 1x1 to 1280
    .globalAvgPool,
    .dense 1280 10 ]
  bnChannels := #[32, 32,16, 96,96,24, ... ]  -- all 49, in order

-- 2. The schedule. ResNet-34's and MobileNetV2's, unchanged.
def efficientnetAdamConfig : VerifiedConfig where
  epochs    := 80
  batchSize := 32

-- 3. The entry point: AdamW, cosine with a 3-epoch warmup, baseLR 1e-3.
--    The rms variant, the EMA decay and the drop rate are env knobs.
def runEfficientNetAdam (argv : List String) : IO Unit := do
  let variant := (← IO.getEnv "LEAN_MLIR_VARIANT").getD "adam"
  let bs := ((← IO.getEnv "LEAN_MLIR_BATCH").bind (·.toNat?)).getD
              efficientnetAdamConfig.batchSize
  let sched := enetRmsSchedule
  let rms := variant.contains "rms"
  let baseLR := match (← IO.getEnv "LEAN_MLIR_BASE_LR_U").bind (·.toNat?) with
    | some u => u.toFloat * 1e-6
    | none   => if rms then sched.lr * bs.toFloat / 256.0 else 0.001
  let emaDecay := match (← IO.getEnv "LEAN_MLIR_EMA_DECAY_U").bind (·.toNat?) with
    | some u => u.toFloat * 1e-6
    | none   => 0.9999
  let dropNet := match (← IO.getEnv "LEAN_MLIR_DROP_RATE_U").bind (·.toNat?) with
    | some u => { efficientnetVerified.toNet with
                  dropKeeps := efficientnetVerified.dropKeepsAt (u.toFloat * 1e-6) }
    | none   => efficientnetVerified.toNet
  dropNet.trainAdamSched { efficientnetAdamConfig with batchSize := bs }
    (argv.head?.getD "data") baseLR 0.9 0.999 3 variant
    (if rms then sched.decayRate else 0.0) sched.decayEpochs emaDecay
def main (argv : List String) : IO Unit := runEfficientNetAdam argv

.mbConvSENB ic mid oc r k is one block per line, as Chapter 6’s constructor was, with mid the expanded width; k is the depthwise kernel, mixing \(3 \times 3\) and \(5 \times 5\) by stage, and the expand \(1 \times 1\) is skipped exactly when mid == ic, which is stage 1 and its \(t=1\).

r is the squeeze width, and it is sized off ic. Every block in the listing has r == ic/4: 32 gives 8, 16 gives 4, 192 gives 48. That is the correctness note this chapter returns to in §7.5 from the other direction, and it is worth reading here first, because sizing the squeeze off the \(6\times \)-larger mid instead is the difference between B0’s canonical 5.29M parameters and an inflated 8.4M. The spec is where that choice is made once and visibly.

The arm, and where the SE gate lives in it.

  | mbConvSENB ic mid oc r k =>                -- Ch. 7
    (if mid != ic then #[(#[mid,ic,1,1],0),(#[mid],1),(#[mid],2)] else #[])
    ++ #[(#[mid,1,k,k],0),(#[mid],1),(#[mid],2),     -- depthwise kxk + BN
         (#[mid,r],0),(#[r],2),(#[r,mid],0),(#[mid],2),   -- SE: squeeze, excite
         (#[oc,mid,1,1],0),(#[oc],1),(#[oc],2)]     -- project 1x1 + BN

Same mid != ic conditional as Chapter 6’s, with four tensors inserted in the middle: [mid,r] down to the squeeze width and [r,mid] back up, each with a bias. Those two are the only convs in the block that keep their biases, because they are not BN-followed — the gate ends in a sigmoid, not a normalizer. That asymmetry is visible in the arm and nowhere else in the spec.

There is no stride argument, and that is deliberate. B0’s downsampling pattern is fixed by the architecture rather than chosen per block, so it lives in the renderer, which emits b1 as no-expand, b2, b4, b6 and b12 as strided, and the remaining nine as residual. The spec’s job is the parameter layout, and efficientnetVerified.toSpecs is kernel-#guarded against an independently audited hand-list.

The BatchNorm list has 49 entries, not the 50 a per-block count gives. Every MBConv normalizes after its expand, depthwise and project convs, which would be \(16 \times 3 + 2\) for the stem and head, or 50. Stage 1 has \(t=1\) and therefore no expand conv, so it contributes two entries rather than three, and the real count is 49. bnChannels exists to get that right: the running-stat region is positional, and an off-by-one misaligns every frozen statistic downstream of it at evaluation time.

The squeeze-excite biases stay. NB drops the bias from every convolution a BatchNorm follows, because the BatchNorm’s own shift absorbs it. SE’s two \(1 \times 1\)s are followed by a sigmoid gate rather than a BatchNorm, so nothing absorbs theirs and they are kept. The rule is about what comes after a convolution, not about the convolution.

Every SE block multiplies the input by the sigmoid output of its two-layer MLP: that is exactly Theorem 96 composed with Theorem 13 and the dense VJP (Theorem 19). No new math in this chapter beyond the SE block itself.

Results

§7.1 has the run: 80 epochs at about 30 seconds each, forty-one minutes, \(\mathbf{90.04\% }\) top-1 and \(\mathbf{98.29\% }\) top-5, with the loss settling on the same \({\sim }0.50\) label-smoothing floor as ResNet-34 and MobileNetV2. Table 7-1 is the three-chapter comparison — same dataset, same recipe, same eighty epochs, every row a measurement on the verified XLA path, on one RTX 4060 Ti.

ResNet-34’s accuracies are the five-seed means of §5.5.2 (runs/2026-09-12-r34-ablation-fp32-seeds/) and its times are one of those seeds’ (§5.1); the other two rows are one run each.

• EfficientNet-B0 is level with ResNet-34 on top-1, at \(90.04\% \) against \(89.99 \pm 0.32\), and eight tenths of a point over MobileNetV2’s \(89.25\% \). It gets there with \(5.3\times \) fewer parameters than ResNet-34 and in six fewer minutes. MobileNetV2 keeps top-5 by four tenths, which is the one place the larger model does not win.

• SE and Swish buy \(+0.71\) top-1 over MobileNetV2 for \(1.80\times \) the parameters (4.02M against 2.24M). That is a real return rather than a bargain: the cost is the two-layer MLP inside every one of the sixteen MBConv blocks. Whether \(1.8\times \) the parameters is worth seven tenths of a point is a judgement, but it is a much better trade than it looks on the ImageNet leaderboard, where B0 is usually compared against far larger models.

• 1,317 KB of MLIR is the largest render in the book, ahead of MobileNetV2’s 1,071 KB and ResNet-34’s 754 KB. The extra ops are SE’s: each block emits a global-average-pool, two denses, a sigmoid and an elementwise multiply, and then the product-rule backward for all of it. Sixteen blocks of that is a lot of StableHLO.

• 103 ms per step sits between the two, slower than MobileNetV2’s 90 ms because of exactly those extra ops, and quicker than ResNet-34’s 120 ms by less than the parameter count suggests, for the reason Chapter 6 gave: the core is a depthwise sandwich, and depthwise is memory-bound.

• The lift is not from SE alone. EfficientNet’s paper thesis is compound scaling, tuning depth, width and resolution together with one coefficient to reach the ImageNet state of the art. At the B0 baseline shown here SE gives a modest bump, and the real wins arrive at B3–B7 where the scaling compounds. This example sits at B0 because the dataset is 10-class Imagenette rather than 1000-class ImageNet, and the return on scaling saturates much earlier.

7.4 MLIR: Squeeze-and-Excitation

What is already proven. Squeeze-excitation is an elementwise product \(\mathrm{se}(x) = x \odot g(x)\), where the per-channel gate \(g(x)\) is itself a small sub-network of the same input \(x\): global-average-pool, dense, swish, dense, sigmoid, broadcast. Because the gate depends on \(x\), the VJP is a product rule, not a plain scaling: seBlockHasVJP proves

the gate-weighted cotangent plus the gate’s own backward applied to the co-input \(x\odot dy\). seBlockFullHasVJP supplies the concrete gate (the GAP–dense–swish–dense–sigmoid chain) and mbconvBodyHasVJP composes the whole MBConv.

The gap and how we close it. The emitted graph realizes seBlockHasVJP’s product-rule backward in full, both terms, with each piece of the gate carrying the swish, sigmoid and dense bridges of the earlier chapters. Here it is in verified_mlir/efficientnet_adam_train_step.mlir at the sixteenth MBConv block (1152 mid-channels, 48 in the squeeze, \(7\times 7\), batch 32), with the swish backward between the two dense backwards elided (SSA names the artifact’s, default precision dropped, lines wrapped):

%v1845 = stablehlo.logistic %v1844 : tensor<32x1152xf32>
%v1847 = stablehlo.broadcast_in_dim %v1845, dims = [0, 1]
      : (tensor<32x1152xf32>) -> tensor<32x1152x7x7xf32>
%v1848 = stablehlo.multiply %v1847, %v1846 : tensor<32x1152x7x7xf32>
%v1849 = stablehlo.multiply %v1832, %v1846 : tensor<32x1152x7x7xf32>
%v1850 = stablehlo.reduce(%v1849 init: %v1833)
      applies stablehlo.add across dimensions = [2, 3]
      : (tensor<32x1152x7x7xf32>, tensor<f32>) -> tensor<32x1152xf32>
%v1851 = stablehlo.constant dense<1.0> : tensor<32x1152xf32>
%v1852 = stablehlo.subtract %v1851, %v1845 : tensor<32x1152xf32>
%v1853 = stablehlo.multiply %v1845, %v1852 : tensor<32x1152xf32>
%v1854 = stablehlo.multiply %v1850, %v1853 : tensor<32x1152xf32>
%v1855 = stablehlo.dot_general %v1854, %b16zW2,
      contracting_dims = [1] x [1] : (tensor<32x1152xf32>,
      tensor<48x1152xf32>) -> tensor<32x48xf32>
// ...
%v1863 = stablehlo.dot_general %v1862, %b16zW1,
      contracting_dims = [1] x [1] : (tensor<32x48xf32>,
      tensor<1152x48xf32>) -> tensor<32x1152xf32>
%v1864 = stablehlo.constant dense<49.0> : tensor<32x1152xf32>
%v1865 = stablehlo.divide %v1863, %v1864 : tensor<32x1152xf32>
%v1866 = stablehlo.broadcast_in_dim %v1865, dims = [0, 1]
      : (tensor<32x1152xf32>) -> tensor<32x1152x7x7xf32>
%v1867 = stablehlo.add %v1848, %v1866 : tensor<32x1152x7x7xf32>

Read it against the product rule, with %v1846 the cotangent \(dy\) into the block and %v1832 its input \(x\):

• %v1845 is the gate \(g(x)\), the logistic of the excite MLP’s output, recomputed; %v1848 is \(g(x)\odot dy\) after the broadcast %v1847, the term a naive generator would stop at.

• %v1849 is \(x \odot dy\) and %v1850 reduces it over the spatial axes. That reduction is the adjoint of the gate’s broadcast, so it is the cotangent on the per-channel gate values. %v1851–%v1854 multiply by \(\sigma (1-\sigma )\), the sigmoid’s derivative, %v1855 and %v1863 are the two dense backwards (each contracting the output axis of its weight), and %v1865 divides by \(49 = 7 \cdot 7\), the adjoint of the global-average-pool, before %v1866 broadcasts the result back over the map.

• %v1867 = %v1848 + %v1866 adds the two.

That final add is the product rule: the gradient that passes straight through the channel, plus the gradient that re-tunes the gate, summed. The leaf theorem is seBackBatched_faithful: the node denotes seBHasVJP’s backward at the saved activation \(x\). SE is not linear, so unlike the convolution leaves this one threads each example’s forward activation through the backward, and the batched statement is built from the per-example se_back_bridge row by row. Theorem 99 carries it through the sixteen blocks.

Caveats.

• The SE fan-in carries no corner condition. Swish and sigmoid are smooth everywhere, and the only smooth-point caveats in an MBConv come from its BatchNorms (\(\epsilon {\gt} 0\)).

• Shape. The listing is one block of the trained net; the theorem is stated at every channel count and squeeze width.

7.5 ImageNet recipe

As with ResNet-34 (Chapter 5) and MobileNet-V2, the B0 backbone scales to full 1000-class ImageNet unchanged. The MBConv stack stays put, the head widens to 1000, and the schedule grows. As there, the recipe runs on both paths — the reference first, then the verified path, then the two side by side.

The reference (JAX)

The reference trainer is jax/MainEfficientNetImagenet.lean; its spec, and its recipe, as the run below trained it:

-- Same MBConv (SE + Swish) backbone, 1000 classes instead of 10.
def efficientNetB0Imagenet : NetSpec where
  name   := "EfficientNet-B0 (ImageNet, bf16)"
  imageH := 224
  imageW := 224
  convBnAct := .swish   -- stem and head too, not just the MBConv interiors
  layers := [
    .convBn 3 32 3 2 .same,
    .mbConv  32  16 1 3 1 1 true,
    .mbConv  16  24 6 3 2 2 true,
    .mbConv  24  40 6 5 2 2 true,
    .mbConv  40  80 6 3 2 3 true,
    .mbConv  80 112 6 5 1 3 true,
    .mbConv 112 192 6 5 2 4 true,
    .mbConv 192 320 6 3 1 1 true,
    .convBn 320 1280 1 1 .same,
    .globalAvgPool,
    .dense 1280 1000 .identity      -- 1000-class head
  ]

-- TF TPU EfficientNet's own code, row for row; the table below names each source.
def efficientNetB0ImagenetConfig : TrainConfig where
  learningRate   := 0.016   -- 0.256 at batch 4096, scaled to 256
  batchSize      := 256
  epochs         := 80
  optimizer      := .rmsprop  -- EfficientNet's original optimizer
  momentum       := 0.9       -- μ for the RMSprop momentum buffer
  rmspropDecay   := 0.9       -- ρ, the running mean-square decay
  rmspropEps     := 1e-3      -- EfficientNet uses ε=1e-3
  gradClipNorm   := 0.0       -- off: the paper uses none
  weightDecay    := 1e-5
  wdExcludeNormBias := true   -- no decay on BN γ/β or biases
  cosineDecay      := false   -- the paper's exponential decay instead
  expLRDecayRate   := 0.97    -- ×0.97 every 2.4 epochs
  expLRDecayEpochs := 2.4
  expLRStaircase   := true    -- floored on the global step
  dropout          := 0.2     -- EfficientNet-B0 classifier dropout
  warmupEpochs   := 5
  augment        := true
  useAutoAugment := true     -- full AutoAugment ImageNet policy
  autoAugmentV0  := true     -- TF EfficientNet's policy v0 (timm `v0`)
  augBicubic     := false   -- TF's geometry is nearest, not PIL bicubic
  cropFallbackCenter := true -- a whole-image draw falls back to the centre crop
  labelSmoothing := 0.1
  bf16           := true
  bf16Conv       := true    -- reaches the MBConv expand/depthwise/project
  f32StemHead    := true    -- classifier head in fp32, as every verified render
  useEMA         := true     -- weight averaging at 0.9999, BN buffers shadowed too
  dropPath       := 0.2      -- stochastic depth, the drop-connect rate
  dropPathOverN  := true     -- the ramp 0.2 · i/16, not the default i/15
  runningBN      := true     -- paper-faithful eval: running BN stats
  bnEps          := 1e-3     -- TF's BN ε
  -- ── The rest of TF TPU EfficientNet's own code, set row by row ──
  rmspropTf      := true     -- TF1 ApplyRMSProp folds lr into the momentum buffer
  wdDecayBiases  := true     -- weight decay over every non-BN variable: biases too
  initHeadTfUniform := true  -- U(±1/√1000) on the classifier, zero bias
  autoAugmentTf  := true     -- the v0 table op for op, TF's fills and shifts
  augNearest     := true     -- contrib_image transforms interpolate NEAREST
  randAugmentTranslateConst := 250.0  -- translate_const 250: (m/10)·250 px
  resizeLegacyTf1 := true    -- TF1 resize_bicubic: no antialias
  evalCropTf     := true     -- crop first, 224/256 of the shorter side …
  evalCropTfEfficientNet := true  -- … with TF's offsets, then bicubic to 224
  evalF32        := true     -- the logged number through an f32 forward


/-- Paper-faithful full run: identical recipe at the 350-epoch schedule, at the
    paper's real LR 0.016 (the TF-RMSProp form trains stably at it —
    no lowered LR, no grad clip). Selected with the `full` recipe arg. -/
def efficientNetB0ImagenetConfigFull : TrainConfig :=
  { efficientNetB0ImagenetConfig with epochs := 350 }

The number this section reports is the full tier, the 350-epoch one, from the run on this listing. The recipe is not the paper’s sentence but TensorFlow’s own EfficientNet code read row by row — main.py, utils.py, efficientnet_model.py, preprocessing.py and autoaugment.py — and Table 7-2 names the source each value was read off.

Two details are specific to EfficientNet. First, bf16Conv casts the MBConv block’s heavy convolutions (the \(1\times 1\) expand, the depthwise, the \(1\times 1\) project) but deliberately leaves the squeeze-excitation \(1\times 1\)s in fp32: SE acts on a globally-pooled \(1\times 1\) tensor where there is no throughput to win, and its sigmoid gate is precision-sensitive. The payoff is the same \(\sim 2\times \) MBConv-block speedup MobileNet sees, from exactly the convolutions worth casting.

Second, the SE bottleneck is sized off the block’s input channels (\(c/4\)), not the \(6\times \)-larger expanded width the squeeze sits behind. That distinction is most of the difference between a \(5.3\)M-parameter B0 and an \(8.4\)M one, and it is what makes the spec in §7.3 derive the canonical \(\mathbf{5.29}\)M.

Augmentation. The base is Inception-style random-resized-crop through TensorFlow’s sampler (10–100% of the image area at a 3/4–4/3 aspect ratio, a whole-image draw falling back to the centre crop) resized to \(224\times 224\) by TF 1’s bicubic kernel, plus a random horizontal flip, with label smoothing \(0.1\) in the loss; validation is TensorFlow’s _decode_and_center_crop, \(224/256\) of the shorter side then the same bicubic. On top of that the recipe turns on AutoAugment at policy v0 as autoaugment.py applies it: the table TensorFlow’s EfficientNet trains with, geometric operations included, at nearest-neighbour resampling with fill \(128\) and TranslateY in pixels. The heavier pack (Mixup, CutMix, RandAugment, Random Erasing) is wired into the reference trainer, emitted into the generated tf.data pipeline with each knob gated on a config flag, and left off here because the 2019 recipe reaches for AutoAugment instead. Table 7-3 is the run.

The run reached \(\mathbf{77.17\% }\) top-1 / \(\mathbf{93.39\% }\) top-5 on the full 50,000-image validation split against EfficientNet-B0’s published \(77.1\% \, /\, 93.3\% \) — a match on the number. In four details of the recipe the obvious implementation is not the one B0 was trained with.

RMSProp in TensorFlow’s form. B0’s optimizer is RMSProp (\(\rho =0.9\), \(\mu =0.9\), \(\varepsilon =10^{-3}\)) in TensorFlow’s variant of it, timm’s RMSpropTF: the update divides by \(\sqrt{\mathrm{sq}+\varepsilon }\), with \(\varepsilon \) inside the square root, and the mean-square accumulator initializes to \(1\). The textbook form puts \(\varepsilon \) outside (\(\sqrt{\mathrm{sq}}+\varepsilon \)) and starts the accumulator at \(0\); together those inflate the effective step by up to \(\sim \)30\(\times \) while the running mean-square is small. At B0’s scale that is enough to erode top-1 over training rather than diverge outright, which makes it read as a learning-rate sensitivity. MobileNet-V2 is indifferent to the choice only because its \(\varepsilon =1.0\) swamps the misplacement. In TF form B0 trains stably at the paper’s \(\mathrm{lr}=0.016\) with no gradient clipping and no lowered peak. TF 1’s ApplyRMSProp also folds the learning rate into the momentum buffer rather than applying it on the way out, and the two forms part by a little at every step of the staircase; the trainer carries TF 1’s (Theorem 87).

The EMA shadows the BatchNorm buffers as well as the weights. The paper’s weight EMA (decay \(0.9999\)) averages the parameters (Theorem 88); evaluation then has to pair those averaged weights with equally-lagged normalization statistics, so the trainer carries a second shadow over the BN running mean and variance (ema_bn) and evaluates the pair. Averaged weights against live BN statistics is a mismatch large enough to put validation loss at \(10^{8}\). LayerNorm carries no running buffers, so ConvNeXt and ViT run the same EMA without needing this.

The EMA decay is warmup-corrected. A flat \(0.9999\) from step \(0\) moves the shadow one ten-thousandth of the way toward the live weights per step, so for the first several thousand steps the evaluated weights are mostly the random initialization while the network trains normally underneath. The decay is therefore \(\min (0.9999,\, (1+t)/(10+t))\): it tracks the live weights early and relaxes to \(0.9999\) once there is something worth averaging. It costs nothing at convergence, and the opening epochs measure the model rather than the initialization.

Swish throughout, stem and head included. B0 is SiLU/swish everywhere, not only inside the MBConv blocks. The two .convBn layers that bracket the stack — the \(3\times 3\) stride-2 stem and the \(1\times 1\) head into the classifier — set convBnAct to swish rather than taking its .relu default. Both sit where the entire signal passes, so two activations decide what the whole network sees on the way in and on the way out.

The 350-epoch run. The full tier reached \(\mathbf{77.17\% }\) top-1 / \(\mathbf{93.39\% }\) top-5 against EfficientNet-B0’s paper \(77.1\% \, /\, 93.3\% \): level on top-1 and within \(0.1\) on top-5. The run is from scratch at the paper’s optimizer, learning rate, EMA and activation, with RMSProp at \(0.016\), AutoAugment, stochastic depth and the exponential staircase, and it departs from TensorFlow’s code in no line the listing names. It ran \(42.5\) hr on four RTX 3060 at \(7.3\) minutes per epoch, with no restart. The headline figure is the EMA weights scored over all 50,000 validation images through an fp32 forward; the same checkpoint’s live weights score \(76.63\% \) top-1, so the weight average is worth \(+0.54\) points. The last twenty-five epochs average \(77.18\, /\, 93.41\), so the final number is the tail’s mean, not the top of a rise. Figure 7-3 is the run, epoch by epoch.

The final \(77.17\% \, /\, 93.39\% \) matches the paper’s \(77.1\% \, /\, 93.3\% \). The flattening past epoch \(\sim \)300 is the exponential schedule running out of learning rate (\(0.016 \to 1.9\times 10^{-4}\)), not saturation of the architecture.

The verified path (PJRT)

The verified trainer trains the same network through the proof-rendered StableHLO that §7.1 ran at Imagenette scale. It is apps/imagenette/MainEfficientNetImagenet.lean, built as lake build efficientnet-imagenet-verified:

-- Same MBConv + SE backbone, 1000 classes instead of 10.
def efficientnetImagenetVerified : VerifiedNetSpec where
  name       := "EfficientNet-B0 (ImageNet-1k)"
  slug       := "efficientnetin"
  data       := .imagenet
  shimScript := "generated_efficientnet_b0_imagenet_shim.py"
  layers     := [ ... the Imagenette stack, 1000-class head ... ]

-- 350 epochs at 64 PER DEVICE; four replicas give the global 256.
-- 350 because that is the schedule the reference number above was
-- measured on.
def efficientnetImagenetConfig : VerifiedConfig where
  epochs    := 350
  batchSize := 64

def main (argv : List String) : IO Unit :=
  efficientnetImagenetVerified.toNet.trainAdamSched
    efficientnetImagenetConfig (argv.head?.getD "data") ... variant

slug names verified_mlir/efficientnetin_<variant>_train_step.mlir, the rendered graph, and shimScript names this net’s own augmentation shim, which matters more here than it does for MobileNetV2 because EfficientNet’s reference config turns on the full AutoAugment policy and streaming another net’s shim would drop it. The head is the only parameter shape that moves, from \(1280 \times 10\) to \(1280 \times 1000\), which takes the count from 4,020,358 to 5,288,548. That second number is B0’s canonical \(5.29\)M, the count the correctly-sized squeeze restores.

The recipe, on both paths. RMSProp \(\rho \, 0.9\) / \(\mu \, 0.9\) / \(\varepsilon \, 10^{-3}\), lr \(0.256\) at batch \(4096\) (so \(0.016\) at \(256\)), weight decay \(10^{-5}\), the exponential staircase \(\times 0.97\) every \(2.4\) epochs after a five-epoch warmup, dropout \(0.2\), drop-connect \(0.2\) ramped as \(0.2 \cdot i/16\), AutoAugment v0, label smoothing \(0.1\), EMA \(0.9999\) over the weights and the BatchNorm buffers, BN momentum \(0.99\) and \(\varepsilon \, 10^{-3}\), 350 epochs. All of it is carried on both paths, including every place TensorFlow’s code differs from a naive port: \(\varepsilon \) inside the square root (\(g/\sqrt{s+\varepsilon }\), not \(g/(\sqrt{s}+\varepsilon )\)) with the mean-square accumulator started at \(1.0\); the learning rate folded into the momentum buffer; weight decay off BatchNorm’s \(\gamma \) and \(\beta \) alone, so the biases decay; SAME padding on the four strided depthwise convolutions; TensorFlow’s crop sampler with its centre-crop fallback and TF 1’s bicubic resize at train and eval; the classifier initialized \(U(\pm 1/\sqrt{1000})\) with the depthwise and squeeze-excite fan rules; an fp32 classifier head and an fp32 evaluation forward; and BatchNorm synchronised over the four replicas.

Two of those rows are easy to get wrong, and the parity gate measures both. BatchNorm’s \(\varepsilon \) is \(10^{-3}\), the TF value, on both paths: bnEps on the reference, and the eps0001 render with its efficientnetin_fwd_eval_eps0001 scorer on the verified path. And the four strided depthwise convolutions pad as TensorFlow’s SAME does, the odd pixel of padding on the bottom and right; a symmetric pad splits it evenly, so at stride two the depthwise window sits one pixel off, and measured against tf_efficientnet_b0 on shared weights (scripts/parity/enet_timm_parity.py) that one pixel is \(0.37\) of the logit scale at training resolution and \(0.048\) at evaluation. The reference’s mbconv_block passes ’SAME’ to its depthwise, and the verified render emits the odd-phase strided depthwise at its four sites (depthwiseStride2FlatXla, the stem’s convention: a symmetric stride-one convolution read at the odd positions), with the block, step and data-parallel ties stated on that form; the gate pins tf_efficientnet_b0 itself and shows the symmetric pad red.

The activation agrees on both sides. The verified render is swish throughout — \(178\) stablehlo.logistic and not one stablehlo.maximum — and the reference sets convBnAct to swish on the stem and the \(1\times 1\) head. The two paths describe the same network.

The two paths, side by side

Table 7-4 is the two paths, row by row.

The verified column is the emarmstfdp64dropdowneps0001bf16 render, the one lake run enet-default-4gpu trains, with synchronised BatchNorm. The reference’s in-training evaluation reads \(77.17\% \) at epoch 350 and its saved state rescored through the same fp32 forward counts the same \(38{,}584\) images, so the two numbers are on the same footing; the verified run’s \(77.21\% \) is \(38{,}604\). The verified run’s \(65.1\) h include six planned refreshes of its loader feed, a minute or so each, and one slow stretch of eight epochs behind a degraded loader, about fifteen minutes in all; no restart.

The two paths tie, \(+0.04\) on top-1 (20 images) and \(+0.07\) on top-5, a tenth of the 95% Wilson half-width (\(\pm 0.37\) at \(77\% \) on \(50{,}000\) images) at one seed each, and the curves sit on each other the whole way: the training losses agree to within \(0.007\) in every 50-epoch window, with the verified run’s a little lower in each, and the sampled accuracies never part by more than \(0.30\) at any epoch in the table. Paired on the same images, the two checkpoints disagree on \(4{,}400\), \(2{,}210\) against \(2{,}190\), and McNemar’s exact test gives \(p = 0.77\): they cannot be told apart on this validation set. The live weights, scored beside the averaged ones, tell the same story at \(p = 0.41\), and put the weight average at \(+0.54\) on the reference and \(+0.70\) on the verified path.

The verified curve runs a hair ahead. It leads on top-1 at \(284\) of the \(350\) epochs, by \(+0.02\) on average over the first fifty, by as much as \(+0.23\) over epochs 151–200, and by \(+0.05\) over the last fifty; the mean of epochs 326–350 is \(77.23\) against \(77.18\). None of it is outside one seed’s noise. Table 7-5 is the selected epochs and Figure 7-4 every ten.

The two curves start together as well, \(12.54\) against \(13.07\) at epoch 1, with the initializations and the seeding of the BatchNorm running statistics held equal.

What the EfficientNet recipe changes from the 2018 recipe

The ResNet-34 trainer of Chapter 5 is the 2018-era recipe: SGD with momentum plus the “bag of tricks” polish (cosine schedule, warmup, label smoothing, random-resized-crop). The config above is the 2019 EfficientNet recipe (Tan & Le). Chapter 5’s RSB-A3 side quest shows what 2021 does to the 2018 baseline, and Chapter 8’s recipe table is 2022’s version. Table 7-6 is the 2019 rung of the same ladder. The trainAdamSched entry point is identical in both columns, and every row is a one-line change to the VerifiedNetSpec, the VerifiedConfig, or the arguments beside them.

What carries over unchanged: softmax cross-entropy with label smoothing \(0.1\), the 5-epoch warmup, random-resized-crop + horizontal flip as the base augmentation, batch 256, train and eval both at \(224\), bf16 convolution (the SE \(1\times 1\)s and the classifier head in fp32), and no gradient clipping, which is unlike every later recipe in this ladder.

The two quietest rows carry the chapter’s two debugging lessons, both told in full above. “RMSProp (TF form)” is doing real work: \(\varepsilon \) inside the square root and a mean-square accumulator that starts at \(1\) are the difference between training stably at the paper’s \(0.016\) and eroding top-1 over the run. And the EMA row only evaluates correctly because the BN running statistics are shadowed alongside the weights (ema_bn); averaged weights against live BN statistics is an eval blow-up, a failure mode specific to BatchNorm nets trained with weight averaging.

The through-line of the 2018\(\to \)2019 shift is “smaller, longer, more automated”: a quarter of the parameters trained for \(4\times \) the epochs, with a learned augmentation policy, stochastic depth, and averaged weights supplying the regularization. It is also a reminder that “the paper recipe” includes the optimizer’s exact semantics, not just its hyperparameters.

7.6 Recap

Three operations and the end of a hypothesis. The squeeze-and-excitation gate multiplies a feature map by a gain computed from its own global average, so its backward is the product rule applied to a network, the multiplicative fan-in that completes the three composition rules. Swish and the sigmoid inside the gate are smooth, and that is what Theorem 97 buys: the full B0, sixteen MBConv blocks, with its backward proved at every input, where every network before it had to name a point and assume the activations stayed off their corners. The step is tied (Theorem 99) and the data-parallel twin holds.

Three follow-ups extended the family in orthogonal directions, and none needs a new primitive. Noisy Student (Xie et al. 2019) is semi-supervised distillation: a teacher trained on labeled ImageNet pseudo-labels JFT-300M, a larger student trains on both under heavy augmentation (RandAugment, dropout, stochastic depth, the “noise” of the title), and the loop repeats, which took EfficientNet-L2 to \(88.4\% \) top-1. EfficientDet (Tan et al. 2020) puts the backbone under a bidirectional feature pyramid, BiFPN, which fuses each resolution with its coarser and finer neighbours by weighted sums, and a small per-anchor MLP predicts class and box at every pyramid level. EfficientNetV2 (Tan & Le 2021) replaces the early stages’ expand-\(1\times 1\) plus depthwise \(k\times k\) with one regular \(k\times k\) convolution, Fused MBConv: more FLOPs, less wall clock, because depthwise convolution is cheap in arithmetic and expensive in memory traffic. The fused block is already in the kit as .fusedMbConvNB, MobileNetV4’s stage 0 (§6.6), so V2 is a spec away.

The chapter’s other lesson is that a recipe is a specification, not a list of hyperparameters: RMSProp in TensorFlow’s form, the EMA shadowing the BatchNorm buffers, the warmup-corrected decay, each a detail that moves the number. The two paths land close and not on top of each other, and the section says which rows are the suspects. The next chapter takes the recipe idea to its conclusion, a ResNet modernized one transformer-era choice at a time, and brings the two operations that era runs on, LayerNorm and GELU.