Efficient ConvNets: Depthwise Separability, Mobile Architectures, and Re-parameterization

Where convnets still win

Most deployed convnets run on phones, cameras, cars, and embedded boards: budgets in milliseconds and milliwatts. Three ideas organize the efficient regime:

  • Factorize: depthwise-separable convolutions (MobileNet).
  • Scale and search: let optimization allocate width, depth, resolution (EfficientNet, MobileNetV4).
  • Re-parameterize: train a multi-branch net, ship an algebraically identical single-branch one (RepVGG).

Depthwise separability, recalled

A dense k \times k convolution factorizes into a per-channel spatial filter plus a 1 \times 1 channel mixture, at fraction

\frac{1}{c_\textrm{o}} + \frac{1}{k^2}

of the dense cost: roughly 1/9 for 3 \times 3 kernels.

The depthwise stage filters each channel; the pointwise stage mixes channels; BN and ReLU follow each.

The block in code

MobileNet is this block, stacked; stride-2 depthwise convolutions downsample, so there is no pooling:

class DWSBlock(nnx.Module):
    """Depthwise 3x3 and pointwise 1x1 convolutions, each with BN and ReLU."""
    def __init__(self, in_channels, out_channels, strides=(1, 1), rngs=None):
        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
        he = nnx.initializers.he_normal()
        self.depthwise = nnx.Conv(
            in_channels, in_channels, kernel_size=(3, 3), strides=strides,
            padding='same', feature_group_count=in_channels, use_bias=False,
            kernel_init=he, rngs=rngs)
        self.bn1 = nnx.BatchNorm(in_channels, momentum=0.9, rngs=rngs)
        self.pointwise = nnx.Conv(in_channels, out_channels,
                                  kernel_size=(1, 1), use_bias=False,
                                  kernel_init=he, rngs=rngs)
        self.bn2 = nnx.BatchNorm(out_channels, momentum=0.9, rngs=rngs)

    def __call__(self, X):
        X = nnx.relu(self.bn1(self.depthwise(X)))
        return nnx.relu(self.bn2(self.pointwise(X)))

Mini-MobileNet

Like VGG, the network is a list of block parameters, here (channels, stride) pairs:

class MiniMobileNet(d2l.Classifier):
    def __init__(self, arch=((64, 1), (128, 2), (128, 1), (256, 2),
                             (256, 1), (512, 2), (512, 1)),
                 lr=0.1, num_classes=10, rngs=None):
        super().__init__()
        self.save_hyperparameters(ignore=['rngs'])
        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
        layers = [nnx.Conv(1, 32, kernel_size=(3, 3), strides=(2, 2),
                           padding='same', use_bias=False,
                           kernel_init=nnx.initializers.he_normal(),
                           rngs=rngs),
                  nnx.BatchNorm(32, momentum=0.9, rngs=rngs), nnx.relu]
        c = 32
        for c_out, s in arch:
            layers.append(DWSBlock(c, c_out, (s, s), rngs))
            c = c_out
        layers.extend([lambda x: x.mean(axis=(1, 2)),  # global average pooling
                       nnx.Linear(c, num_classes, rngs=rngs)])
        self.net = nnx.Sequential(*layers)

    def configure_optimizers(self):
        return optax.sgd(self.lr, momentum=0.9)

0.54M trainable parameters; a dense twin at the same widths would need about 4.7M, the 8.8\times gap the cost ratio predicts:

model = MiniMobileNet()
X = jnp.zeros((1, 96, 96, 1))
assert model(X).shape == (1, 10)
sum(p.size for _, p in nnx.state(model, nnx.Param).flat_state())
542474

Matched parameter budget

Same parameter budget, same head, 10 epochs of Fashion-MNIST at 96×96; the dense net must stop at 128 channels where the separable one reaches 512:

for name, (acc, seconds) in runs.items():
    params = sum(p.numel() for p in models[name].parameters())
    print(f'{name}: {params:,} params, '
          f'val acc {acc:.3f}, {seconds:.1f} s/epoch')
mini-MobileNet: 542,474 params, val acc 0.909
VGG-style: 583,594 params, val acc 0.887

Accuracy is comparable across frameworks and seeds; the ordering of the two models is noise. Mini-MobileNet uses 7.7\times fewer multiply-adds (50M vs. 385M), but its epochs are only about 2–3\times faster: FLOPs are not latency (depthwise convolutions are memory-bound).

The inverted bottleneck

ResNet bottlenecks compress around the expensive dense 3×3. Once the 3×3 is depthwise, wide is cheap and MobileNetV2 inverts the block:

  • thin residual stream, 1 \times 1 expansion (6×),
  • depthwise 3×3 in the wide space,
  • linear 1 \times 1 projection back (no ReLU: it would destroy information in the thin space).

The tensors alive across blocks are the thin ones, so activation memory stays small. ConvNeXt reused this exact shape with a 7×7 depthwise convolution.

Scaling and searching

  • EfficientNet: scale depth, width, resolution together (\alpha^\phi, \beta^\phi, \gamma^\phi, one budget knob \phi); base block found by architecture search.
  • EfficientNetV2: search rewards training speed; drops depthwise blocks from early high-resolution stages (memory-bound there) and grows image size during training.
  • MobileNetV4: one universal inverted bottleneck covering four block variants; searched to be near Pareto-optimal across phone CPUs, DSPs, GPUs, and NPUs.

The block is human insight; the allocation is an optimization problem against measured latency.

RepVGG: train branches, ship a stack

Multi-branch nets train better (ResNet’s lesson); plain 3×3 stacks run faster (one big kernel per layer, one live tensor). RepVGG gets both: train with branches, then fuse them exactly.

Train-time: 3×3 + 1×1 + identity, each with BN. Inference: one 3×3 convolution, same function.

The train-time block

class RepVGGBlock(nnx.Module):
    """Train-time block: 3x3, 1x1, and identity branches, each with BN."""
    def __init__(self, num_channels, rngs=None):
        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
        self.conv3 = nnx.Conv(num_channels, num_channels, kernel_size=(3, 3),
                              padding='same', use_bias=False, rngs=rngs)
        self.conv1 = nnx.Conv(num_channels, num_channels, kernel_size=(1, 1),
                              use_bias=False, rngs=rngs)
        self.bn3 = nnx.BatchNorm(num_channels, rngs=rngs)
        self.bn1 = nnx.BatchNorm(num_channels, rngs=rngs)
        self.bn0 = nnx.BatchNorm(num_channels, rngs=rngs)

    def __call__(self, X):
        return nnx.relu(self.bn3(self.conv3(X)) + self.bn1(self.conv1(X))
                        + self.bn0(X))

The fusion algebra

Three pieces of weight algebra, all exact:

  1. Fold BN into the conv (per output channel o):

\hat{\mathbf{W}}_o = \frac{\gamma_o}{\sqrt{\sigma_o^2+\epsilon}}\,\mathbf{W}_o, \qquad \hat{b}_o = \beta_o - \frac{\gamma_o \mu_o}{\sqrt{\sigma_o^2+\epsilon}}.

  1. A 1 \times 1 kernel is a 3 \times 3 kernel padded with zeros; the identity is a 3 \times 3 kernel with a centered 1 where i = o.
  2. Convolution is linear in its kernel: sum the three kernels and the three biases.

Fusing in code

def fuse_bn(W, bn, eps=1e-5):
    scale = bn.scale[...] / jnp.sqrt(bn.var[...] + eps)
    return W * scale, bn.bias[...] - bn.mean[...] * scale

def fuse_block(blk, c):
    pad = ((1, 1), (1, 1), (0, 0), (0, 0))
    W3, b3 = fuse_bn(blk.conv3.kernel[...], blk.bn3)
    W1, b1 = fuse_bn(jnp.pad(blk.conv1.kernel[...], pad), blk.bn1)
    I = jnp.pad(jnp.eye(c).reshape(1, 1, c, c), pad)
    W0, b0 = fuse_bn(I, blk.bn0)
    return W3 + W1 + W0, b3 + b1 + b0

Exact to the last float

Push a few batches through to populate the BN running statistics, then compare the three-branch block against the single fused conv:

c = 8
blk = RepVGGBlock(c)
train_blk = nnx.view(blk, use_running_average=False)
for _ in range(4):  # populate the BN running statistics
    X = jax.random.normal(d2l.get_key(), (32, 16, 16, c))
    train_blk(X)
W_fused, b_fused = fuse_block(blk, c)
X = jax.random.normal(d2l.get_key(), (32, 16, 16, c))
Y = nnx.view(blk, use_running_average=True)(X)
fused = nnx.Conv(c, c, kernel_size=(3, 3), padding='same',
                 rngs=nnx.Rngs(0))
fused.kernel[...] = W_fused
fused.bias[...] = b_fused
Y_fused = nnx.relu(fused(X))
float(jnp.abs(Y - Y_fused).max())
2.86102294921875e-06

Maximum difference: order 10^{-6}, i.e. single-precision roundoff. Same function, one kernel, no BN at inference.

Deployment has the last word

“Mathematically identical” holds in FP32. The fused kernel sums branches of very different scales, so its weight and activation distributions are wide, and naive INT8 quantization collapses: 75.1% → 40.2% top-1 for a fused RepVGG. Quantization-aware re-parameterization fixes it.

The idea industrialized anyway: MobileOne (sub-millisecond phone backbones) and FastViT (re-param conv stages + attention) both ship on consumer devices.

Recap: convnets at the edge, 2026

  • Depthwise separability: \sim k^2 cheaper convolutions, small accuracy cost.
  • Inverted bottleneck: expensive tensors are the transient ones.
  • Compound scaling / search: allocate budget by optimization.
  • Re-parameterization: the trained net and the shipped net differ, connected by exact weight algebra.
  • Cheapest tier: plain depthwise convnets. Flagship phones: conv-attention hybrids. Pure ViT: only with a dedicated NPU.
  • FLOPs are a proxy; hardware keeps score in memory traffic.