Add CPUBone model

[Krystex]

Hi there!

In this PR I want to add the CPUBone by Nottebaum et al. to timm. Disclaimer: I'm not the author of the paper.

I took the reference implementation, cleaned it up and integrated it into the timm library so it can be easily loaded with timm.create_model('cpubone_b0').
There are still some minor cleanups to do, but I think this first draft is ready for a review by the original author.

[rwightman]

@Krystex thx, the timm impl looks in decent shape. Did you verify equivalence to original across all model scales? Does it reproduce imagenet-1k eval numbers?

[altair199797]

Hi,

author here. Looks good, I see no problems. To be sure we'd need to evaluate ImageNet accuracy.
And thanks again for doing that!:D

Best Regards

[rwightman]

@altair199797 @Krystex

So, I took a closer look, reproduced evals in float32 and looked at timm API compat. That checked out. But use with autocast fails, especially bfloat16 which resulted in complete collapse.

I thought it might be an attention logits overflow, similar to EfficientViT-MIT but it's not. So I tasked codex to isolate the cause of the collapse... it found the issue. There's a corresponding float32 forced region that works around, but I feel it is still going to be problematic for downstream use, fine-tune, etc and likely not going to end up being used extensively in CPU or mobile scenarios where further quantization, low prec operation is often the norm...

Codex' summary here, a bit more inflammatory than warranted, but it covers the key points:

---

The architecture permits a genuinely ill-conditioned representation.

The problematic local branch is expand conv + bias + Hardswish → depthwise conv + bias + Hardswish → pointwise conv → BatchNorm, with no normalization before the final projection. In the trained B0/B1 weights, that
projection produces channels whose real variance is around (10^{-7}), then the final BN rescales them by roughly 200–1000×. bf16 cannot preserve the tiny difference between the projection output and BN running mean, so
useful signal is lost before the gain is applied.

That is partly structural and partly learned:

• The branch layout makes this failure mode possible.
• The checkpoint/training dynamics determine how badly it manifests: untreated bf16 top-1 is 30.80% for B0, 0.29% for B1, and 80.44% versus 81.20% for B2, while Nano/B3 are much less affected.
• Recalibrating BN does not solve it; the live activations really do have those tiny variances.

So the fp32 island is a compatibility safeguard, not proof the architecture is naturally AMP-robust. Fixing the underlying design would mean changing normalization/bias placement or training regularization and
retraining the weights. For this PR, I’d only accept the safeguard if its CPU cost is modest; I’m measuring that now because “lightweight CPU model with a forced-fp32 late path” is the important tradeoff.

[rwightman]

This was a recommended architectural change to potential avoid the awkward numerics, though would require retrain to confirm / fix.

  if fused_conv and in_channels < 256:
      local_module = FusedMBConv(
          in_channels=in_channels,
          out_channels=in_channels,
          expand_ratio=expand_ratio,
          use_bias=(True, False),
          kernel_size=2 if small_kernels else 3,
          expand_groups=expand_groups,
          norm_layer=(norm_layer, norm_layer),
          act_layer=(act_layer, None),
      )
  else:
      local_module = MBConv(
          in_channels=in_channels,
          out_channels=in_channels,
          expand_ratio=expand_ratio,
          expand_groups=expand_groups,
          # use_bias=(True, True, False),  # <--- removed
          use_bias=(True, False, False),  # <--- added: depthwise bias is redundant with its new norm
          kernel_size=2 if small_kernels else 3,
          # norm_layer=(None, None, norm_layer),  # <--- removed
          norm_layer=(None, norm_layer, norm_layer),  # <--- added: normalize depthwise output before activation
          act_layer=(act_layer, act_layer, None),
      )