Would optional whole-block scaling be a useful addition to NaFlexViT?

[Alihjt]

Hi @rwightman,

I’ve been experimenting with a small modification to NaFlexViT that I’m calling BlockScale, and wanted to ask whether you’d be interested in it as an optional feature in timm.

The idea is to apply a learnable per-channel multiplier after each complete transformer block, similar to layer scale:

output = block(x) * (1 + delta)

Each block has its own delta vector, initialized from a normal distribution with mean zero and standard deviation 1e-4, so the multiplier starts very close to one. The attention, MLP, and existing LayerScale stay unchanged. The scaling acts on the whole block output, including the residual stream. It adds only 2,304, 4,608, and 9,216 trainable parameters to the Tiny, Small, and Base models, respectively.

For training, I followed the Base recipe published on Hugging Face and also used it for Tiny and Small. Here are the completed seed-42 and 43 results: best EMA ImageNet-1k top-1 over 300 epochs, evaluated at 576 patch tokens.

Seed 42

Model	Original NaFlexViT, non-EMA	BlockScale, non-EMA	┃	Original NaFlexViT, EMA	BlockScale, EMA
Tiny	74.634%	75.078% (+0.444)	┃	74.614%	75.060% (+0.446)
Small	82.750%	83.262% (+0.512)	┃	82.710%	83.266% (+0.556)
Base	84.134%	84.454% (+0.320)	┃	84.358%	84.628% (+0.270)

Seed 43

Model	Original NaFlexViT, non-EMA	BlockScale, non-EMA	┃	Original NaFlexViT, EMA	BlockScale, EMA
Tiny	74.764%	75.068% (+0.304)	┃	74.746%	75.026% (+0.280)
Small	82.970%	83.262% (+0.292)	┃	82.948%	83.248% (+0.300)
Base	83.986%	84.374% (+0.388)	┃	84.284%	84.768% (+0.484)

Would you be open to a small PR for this? I have a working implementation and would be happy to share it and adapt the integration to whatever fits the codebase best.

Thanks for your work on timm!

[rwightman]

might be worth a comparison to SBB ViTs for in1k runs... they didn't use NaFlex in training but the recipe was tuned a bit more...

e.g.
https://gist.github.com/rwightman/fb37c339efd2334177ff99a8083ebbc4#file-train-vit_medium_reg4-4gpu-in1k-yaml
https://gist.github.com/rwightman/fb37c339efd2334177ff99a8083ebbc4#file-train-vit_mediumd_rope_reg1-4gpu-in1k-yaml

https://huggingface.co/timm/vit_mediumd_patch16_rope_reg1_gap_256.sbb_in1k
https://huggingface.co/timm/vit_betwixt_patch16_reg4_gap_256.sbb_in1k

if that holds up could look at how it'd cleanly integrate into multiple vit archs..