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 that predates the transformer-era “attention” by four years and informs every post-2018 vision architecture.

7.1 Run it first

Before the attention mechanism, train the network that uses it. Four commands, and less GPU time than either of the two previous chapters:

lake exe cache get                            # Mathlib oleans, ~30 s
./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.048415 lr=0.000333
  epoch 1: val_acc = 698/3925 = 17.783439%  top5 = 2753/3925 = 70.140127%
Epoch 2/80: loss=1.633202 lr=0.000667
  epoch 2: val_acc = 1356/3925 = 34.547771%  top5 = 2831/3925 = 72.127389%
Epoch 3/80: loss=1.421186 lr=0.001000
  epoch 3: val_acc = 1529/3925 = 38.955414%  top5 = 3113/3925 = 79.312102%
Epoch 4/80: loss=1.281925 lr=0.001000
  epoch 4: val_acc = 2520/3925 = 64.203822%  top5 = 3667/3925 = 93.426752%
Epoch 79/80: loss=0.504021 lr=0.000000
  epoch 79: val_acc = 3535/3925 = 90.063694%  top5 = 3865/3925 = 98.471338%
Epoch 80/80: loss=0.504361 lr=0.000000
  epoch 80: val_acc = 3531/3925 = 89.961783%  top5 = 3864/3925 = 98.445860%
done (trained EfficientNet-B0 adam + cosine/warmup via packed threading).

Eighty epochs at about 30 seconds each, forty and a half minutes in total, and 89.96% top-1 with 98.45% top-5 on Imagenette’s 3,925-image validation split. The best epoch reached \(90.14\% \). That is the highest top-1 in this book’s Imagenette column: ResNet-34 finished at \(89.71\% \) and MobileNetV2 at \(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 less than half the wall-clock.

The first four epochs are worth a glance. Top-1 climbs 17.8, 34.5, 39.0, then jumps to 64.2. The learning rate is still warming up through epoch 3, and the BatchNorm running statistics are averaging over weights that are moving fast, so the evaluation forward is using statistics that no longer describe what is being evaluated. It resolves once the schedule reaches its peak and never recurs.

Squeeze-excitation is the first place the product rule carries real weight. That run did not execute a reimplementation of the network this chapter describes. It executed verified_mlir/efficientnet_adam_train_step.mlir, which is pretty(provenGraph) off the renderer. Every squeeze-excitation gate in it backpropagates by the product rule Theorem 42 proves, every depthwise convolution by Chapter 6’s, every BatchNorm by Chapter 4’s, and every residual skip by Chapter 5’s additive fan-in.

And this fold holds at every input. efficientnetForwardB_full_has_vjp chains the stem, all sixteen MBConv blocks and the head through vjp_comp, batched over \(N\), assuming a positive \(\varepsilon \) at each BatchNorm and nothing else — so every operation in the graph that just trained carries a proved backward, and their composition into the whole network carries one too. What buys that is swish. It is smooth, so the statement quantifies over all inputs, where Chapter 6’s fold had to name a point and assume its thirty-five relu6 activations stayed off the kink. Depth was never the obstacle, the activation was — and every network after this one is smooth.

The three compile lines are the same three the two previous chapters showed. @efficientnet_fwd_eval is a separate forward because BatchNorm evaluates differently than it trains, and running-stats BN: 49 layers, 42016 stat floats is the driver threading the running averages out of the train step and into it. The train step reports 740 outputs, which is what 213 parameter tensors plus their Adam moments plus the batch statistics plus the report-only loss come to.

To run this net alongside the other six at the same scale and recipe, see §7.5 for what changes at ImageNet scale.

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 (Chapter 1). We already proved the VJP of dense (Chapter 2). We already proved the VJP of identity, and an additive fan-in collapses to the same product-rule structure as elementwise. Theorem 42 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. Every post-2018 vision architecture has SE or one of its descendants (CBAM, ECA, GE) bolted in, 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 theorem

7.3 Example: EfficientNet-B0 on Imagenette

Squeeze-and-Excitation by itself is one of the smallest architectural contributions in deep learning. Take a feature map, compute a per-channel scalar by global-average-pooling it, pass those scalars through a tiny two-layer MLP with a sigmoid at the end, multiply those sigmoid outputs back into the original feature map per-channel. That’s it. It’s a learned per-channel reweighting, a particularly simple form of attention that predates the transformer-era use of “attention” by four years (Hu et al. 2018, Squeeze-and- Excitation Networks).

EfficientNet (Tan & Le 2019) bolts SE into every MBConv block in what’s otherwise a MobileNet-V2-shaped architecture. The paper’s actual contribution is compound scaling (tune depth, width, and input resolution together) but the B0 baseline’s architectural novelty over MobileNet V2 is essentially “add SE blocks.”

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. The inset shows one MBConv block, an inverted residual with an SE gate (highlighted) inserted before the project.

The verified trainer

§5.3 is the pattern and Chapter 6 was the first chapter to reuse it. Everything that carries weight there carries the same weight here, and only the layers list and the bnChannels that follows from it are new:

import LeanMlir

-- 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
  nClasses := 10
  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.
def main (argv : List String) : IO Unit :=
  efficientnetVerified.toNet.trainAdamSched
    efficientnetAdamConfig (argv.head?.getD "data")
    0.001 0.9 0.999 3 "adam"

.mbConvSENB ic mid oc r k is one block, not a stage: there is no repeat count, so the sixteen blocks are sixteen lines and the stage boundaries are visible as the places where the input width changes. It takes the expanded width mid rather than a ratio \(t\), which is why the widths read straight off the paper’s table, and k is the depthwise kernel, mixing \(3 \times 3\) and \(5 \times 5\) by stage. 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. So the listing above is not a description of the network that got trained, it is the object the renderer consumed.

Forty-nine BatchNorms, and the list is not three per block. 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 § 42 composed with § 12 and the already-proved dense VJPs from Chapter 2. 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 and a half minutes, \(\mathbf{89.96\% }\) top-1 and \(\mathbf{98.45\% }\) top-5, with the loss settling on the same \({\sim }0.50\) label-smoothing floor as ResNet-34 and MobileNetV2. The three-chapter comparison — same dataset, same recipe, same eighty epochs, and every row now a measurement on the verified XLA path, on one RTX 4060 Ti:

Step time is the epoch wall-clock divided by 295 batches, so it includes each epoch’s validation pass. MLIR sizes are the committed <slug>_adam_train_step.mlir in 1000-byte units.

• EfficientNet-B0 has the best top-1 in the column, at \(89.96\% \) against ResNet-34’s \(89.71\% \) and MobileNetV2’s \(89.25\% \). It gets there with \(5.3\times \) fewer parameters than ResNet-34 and in less than half the time. MobileNetV2 keeps top-5 by two 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,316 KB of MLIR is the largest render in the book, ahead of MobileNetV2’s 1,047 KB and ResNet-34’s 729 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 far quicker than ResNet-34’s 220 ms because the core is still a depthwise sandwich rather than full convolutions.

• 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.

Every post-2018 vision architecture has some variant of SE in it (MobileNet V3, EfficientNetV2, ConvNeXt, RegNetZ) because the accuracy-per-parameter cost is low when you compound it with other scaling tricks. The other architectural contribution EfficientNet shipped is the Swish activation \(\mathrm{swish}(x) = x \cdot \sigma (x)\), which is orthogonal to SE.

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: seBlock_has_vjp proves

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

The gap and how we close it. The emitted graph realizes seBlock_has_vjp’s product-rule backward in full. Both terms are there, and each piece of the gate carries the swish, sigmoid, and dense bridges of the previous chapters. Here is the fan-in, with the excite sub-network’s internal VJP elided to its one comment line (the SE block runs on the MBConv’s four mid-channels):

// se = x * broadcast(gate(x)); gate = squeeze-excite sub-network
// %dse = cotangent into the SE block;  %ggb2 = broadcast(gate)
%gdleft = stablehlo.multiply %ggb2, %dse : tensor<1x4x4x4xf32>
%gxdse = stablehlo.multiply %d_s4, %dse : tensor<1x4x4x4xf32>
%gdgate = stablehlo.reduce(%gxdse init: %sc)
            applies stablehlo.add across dimensions = [2, 3]
            : (tensor<1x4x4x4xf32>, tensor<f32>) -> tensor<1x4xf32>
// ... gate VJP through the excite MLP (sigmoid/dense/swish/dense) ...
%gdgate_sp = stablehlo.broadcast_in_dim %gdsq_d, dims = [0, 1]
            : (tensor<1x4xf32>) -> tensor<1x4x4x4xf32>
%gdds = stablehlo.add %gdleft, %gdgate_sp : tensor<1x4x4x4xf32>

Read it against the product rule:

• %gdleft is \(g(x)\odot dy\), the term a naive generator would stop at.

• %gdgate reduces \(x \odot dy\) 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, and it flows back through the excite MLP and the global-average-pool to %gdgate_sp.

• %gdds = %gdleft + %gdgate_sp 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.

As with attention, the listing shows the new structural move (the product-rule fan-in) and elides the work behind it: the // gate VJP comment stands for the excite MLP’s own backward (sigmoid, dense, swish, dense), each carried by this chapter’s bridges and composed in seBlockFull_has_vjp.

Caveats.

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

• Representative scale (four mid-channels, two in the squeeze).

7.5 ImageNet recipe

[TODO: run with timm tweaks.] The numbers in this section were measured before the trainers moved onto timm’s validation protocol and its antialiased resampler (appendix A). That change moves the training distribution as well as the evaluation, so it needs a re-run rather than a re-score, and none of it is reflected below.

Why the phase-2 trainer? The ImageNet numbers in this book come from the phase-2 (Lean\(\to \)JAX) trainer: at this scale its job is to validate the framework’s logic end to end, and it is the reference the verified path is measured against. Phase-3 verified-IREE codegen led to phase 4 — PJRT, which is the verified path itself. Where a section has both, it reports both.

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. The spec mirrors jax/MainEfficientNetImagenet.lean:

-- Same MBConv (SE + Swish) backbone, 1000 classes instead of 10.
def efficientNetB0Imagenet : NetSpec where
  name   := "EfficientNet-B0 (ImageNet, bf16)"
  imageH := 224
  imageW := 224
  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
  ]

-- The paper's own recipe: RMSProp in TensorFlow's form, exponential
-- LR decay, AutoAugment, stochastic depth, EMA, bf16 conv on.
def efficientNetB0ImagenetConfig : TrainConfig where
  learningRate   := 0.016   -- 0.256@4096, linearly scaled to bs 256
  batchSize      := 256
  epochs         := 80      -- a short check; `Full` below is the 350
                            -- this section reports
  optimizer      := .rmsprop
  momentum       := 0.9     -- mu, for the RMSProp momentum buffer
  rmspropDecay   := 0.9     -- rho, the running mean-square decay
  rmspropEps     := 1e-3
  gradClipNorm   := 0.0     -- OFF: the paper uses none
  weightDecay    := 1e-5    -- small: protects depthwise/SE params
  cosineDecay      := false -- replaced by the paper's exp decay
  expLRDecayRate   := 0.97  -- x0.97 every 2.4 epochs, after warmup
  expLRDecayEpochs := 2.4
  dropout          := 0.2   -- classifier dropout
  warmupEpochs   := 5
  augment        := true
  useAutoAugment := true    -- the full ImageNet policy
  labelSmoothing := 0.1
  bf16           := true
  bf16Conv       := true    -- reaches the MBConv expand/dw/project
  useEMA         := true    -- weight averaging, decay 0.9999
  dropPath       := 0.2     -- stochastic depth (B0 drop-connect rate)
  runningBN      := true    -- eval on running stats, not batch stats

-- The 350-epoch tier is the same recipe, longer.
def efficientNetB0ImagenetConfigFull : TrainConfig :=
  { efficientNetB0ImagenetConfig with epochs := 350 }

The number this section reports is the full recipe, the 350-epoch one. The default at 80 epochs is a shorter validation pass over the identical configuration, useful for checking that a change trains at all before committing a week of GPU time to it.

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. An earlier version of the codegen sized it off the expanded width and inflated B0 from its canonical \(5.3\)M parameters to \(8.4\)M. Sizing the squeeze correctly restores the faithful \(\mathbf{5.29}\)M-parameter B0, which is the count §7.3’s spec derives today.

Augmentation. The base is Inception-style random-resized-crop (sampling 8–100% of the image area at a 3/4–4/3 aspect ratio, resized to \(224\times 224\)) plus a random horizontal flip, with label smoothing \(0.1\) in the loss, and validation center-crops. On top of that the run turns on AutoAugment, the full ImageNet policy including its geometric operations, which is the learned augmentation B0’s recipe calls for. The heavier pack (Mixup, CutMix, RandAugment, Random Erasing) is wired into the phase-2 Lean\(\to \)JAX 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.

(CUDA, bf16, batch 256.) That is the paper recipe end to end, and it reached \(\mathbf{76.80\% }\) top-1 / \(\mathbf{93.26\% }\) top-5 on the full 50,000-image validation split against EfficientNet-B0’s published \(77.1\% \, /\, 93.3\% \). Getting there took two bug fixes, and neither was where we first looked.

RMSProp: a cautionary result that was really a bug. EfficientNet’s native optimizer is RMSProp (\(\rho =0.9\), \(\mu =0.9\), \(\varepsilon =10^{-3}\)), and re-running B0 with it exposed what looked like a learning-rate sensitivity absent from MobileNet-V2’s RMSProp: at the MobileNet-style peak \(\mathrm{lr}=0.045\) the loss diverged by epoch 6, and even at the paper’s linear-scaled \(\mathrm{lr}=0.016\) (\(0.256\) at batch \(4096 \to 0.016\) at batch \(256\)) top-1 eroded, peaking \(\sim \)31% near epoch 4, then decaying to \(\sim \)19% and stalling. The obvious workaround was to lower the peak to \(\sim \)0.01, and for a while we did. But the learning rate was not the culprit: our RMSProp was not TensorFlow’s RMSProp, the variant EfficientNet actually trained with. Two differences inflate the effective step by up to \(\sim \)30\(\times \) when the running mean-square is small, which is exactly the erosion observed. \(\varepsilon \) is added outside the square root (\(\sqrt{\mathrm{sq}}+\varepsilon \)) rather than inside (\(\sqrt{\mathrm{sq}+\varepsilon }\)), and the mean-square accumulator is initialized to \(0\) rather than \(1\). This also explains the MobileNet-V2 contrast: its \(\varepsilon =1.0\) is large enough to swamp the misplacement, so the same bug left it unscathed. MNv2 trained fine on the buggy optimizer while B0 did not. Correcting both to the TensorFlow form (timm’s RMSpropTF) let B0 train stably at the paper’s real \(0.016\) with no gradient clipping and no lowered LR.

A second bug surfaced once the optimizer was fixed: with the paper’s weight EMA (decay \(0.9999\)) on, evaluation blew up and validation loss reached \(10^{8}\), because the EMA shadowed only the weights, and eval then paired those averaged weights with the live BatchNorm running statistics. Shadowing the BN buffers with the same EMA (ema_bn) fixed it. ConvNeXt uses the identical EMA and never hit this: LayerNorm has no running buffers to desynchronize, so the bug is specific to BatchNorm nets trained with weight averaging.

The faithful 350-epoch run. With both fixed, the full tier reached \(\mathbf{76.80\% }\) top-1 / \(\mathbf{93.26\% }\) top-5, against EfficientNet-B0’s paper \(77.1\% \, /\, 93.3\% \): within \(0.3\% \) top-1 and essentially exact on top-5. That tier is the paper recipe end to end, with RMSProp at \(0.016\), AutoAugment, stochastic depth, EMA and exponential LR decay. A faithful, from-scratch reproduction with no hacks: real RMSProp, real LR, real EMA. It ran \(\sim \)55.5 hr on four 4060 Ti (350 epochs plus eleven 30-minute thermal rests, zero AER). In the validation curve below, the deliberate slow start is the weight EMA at \(0.9999\) lagging near-init weights until \(\sim \)epoch 8 before catching up sharply:

EfficientNet-B0 / ImageNet-1k validation accuracy per epoch, faithful 350-epoch RMSProp run (bf16, 4\(\times \) 4060 Ti). The EMA slow start (near-zero through \(\sim \)epoch 7) is the weight average lagging the live weights, not a training failure, and the final \(76.80\% \, /\, 93.26\% \) matches the paper’s \(77.1\% \, /\, 93.3\% \).

Phase 4: the verified trainer

Everything above this point is the phase-2 Lean\(\to \)JAX trainer. The phase-4 peer, which trains the same network through the proof-rendered StableHLO that §7.1 ran at Imagenette scale, exists and builds. 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"
  nClasses   := 1000
  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 phase-2 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 phase-2 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.

Where this differs from the paper. The recipe is 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}\), exponential decay \(\times 0.97\) every \(2.4\) epochs after a five-epoch warmup, dropout \(0.2\), drop-connect \(0.2\), AutoAugment, label smoothing \(0.1\), EMA \(0.9999\), BN momentum \(0.99\), 350 epochs. All of it is carried, in both phases, including the two places the TensorFlow original differs from a naive port: \(\varepsilon \) goes inside the square root (\(g/\sqrt{s+\varepsilon }\), not \(g/(\sqrt{s}+\varepsilon )\)), and the mean-square accumulator starts at \(1.0\) rather than \(0\). Drop-connect is a linear ramp over blocks, as the official implementation does it, and the EMA shadows the BN buffers as well as the weights.

One difference is known and open: batch-norm \(\varepsilon \) is \(10^{-5}\) here against the TF reference’s \(10^{-3}\). _bn(..., eps=1e-5) in the emitted phase-2 code, and \(147\) splat constants of 1.0e-5 in the phase-4 render — so the two phases agree with each other and both differ from the paper. It is unquantified: nobody has run the \(10^{-3}\) arm. Phase 2 scores \(76.80\% \) / \(93.26\% \) against the paper’s \(77.1\) / \(93.3\), and that \(-0.3\) is left unattributed rather than pinned on this.

A second difference was found after that run and is now fixed, so the curve above predates it. B0 is SiLU/swish throughout, stem and head included, and the phase-4 render always was — \(194\) stablehlo.logistic and not one stablehlo.maximum. The phase-2 spec, though, never set convBnAct, which defaults to .relu, so the two .convBn layers that bracket the MBConv stack — the stem and the \(1\times 1\) head — trained with ReLU while every block interior used swish.

What it has not done is run. No phase-4 ImageNet result exists for this network, which is why every accuracy above comes from phase 2. What is measured is the throughput, and the schedule follows from it:

The columns are the two arms of emarmsdp64dropdo — EMA, stochastic depth and classifier dropout — the recipe enet-default-4gpu names. fp32 is the default; the run will be bf16, so read the right-hand column. Set that against the \(\sim \)8.6 minutes per epoch the phase-2 350-epoch row reports on the same four cards: both are bf16 over the same architecture, and the two agree to within \(4\% \).

[TODO: run efficientnet-imagenet-verified.]

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 350-epoch run above is the 2019 EfficientNet recipe (Tan & Le), reproduced end to end. Chapter 5’s RSB-A3 side quest shows what 2021 does to the 2018 baseline, and Chapter 8 closes with 2022’s version. This table 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 matmul and convolution (the SE \(1\times 1\)s 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. “RMSProp (TF form)” is doing real work: as recounted above, TensorFlow’s \(\varepsilon \)-outside-the-square-root and mean-square init of \(1\) are the difference between training stably at the paper’s \(0.016\) and eroding to \(\sim \)19%. 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 Further reading

Three EfficientNet follow-ups extended the family in orthogonal directions:

Noisy Student (Xie et al. 2019). Semi-supervised distillation that pushed EfficientNet-L2 to \(88.4\% \) top-1 on ImageNet, a state-of-the-art mark at the time. The recipe: train a teacher on labeled ImageNet, use the teacher to pseudo-label JFT-300M (\(\sim \)300M unlabeled images), then train a larger student on the combined dataset with heavy augmentation (RandAugment, dropout, stochastic depth), which is the “noise” the title refers to. Iterate.

EfficientDet (Tan et al. 2020). EfficientNet backbone plus BiFPN (bidirectional feature pyramid) for object detection. BiFPN aggregates multi-scale features by fusing each spatial resolution with both its coarser and finer neighbors via weighted-sum cross-scale connections. The detection head is a small per-anchor MLP predicting class plus box coordinates, applied at every BiFPN output level.

EfficientNetV2 (Tan & Le 2021). Fused MBConv replaces the early stages’ expand-\(1\times 1\) plus depthwise \(k\times k\) with a single regular \(k\times k\) convolution — more FLOPs, less wall clock, because depthwise convolution is cheap in arithmetic and expensive in memory traffic, and at low resolution the traffic dominates. The fused block is already in the kit as .fusedMbConvNB (§6.6) — it is MobileNetV4’s stage 0 — so V2 needs no new primitive, only a spec.