[CoAtNet] How many heads should be used in SelfAttention?

[xskxzr]

I notice that your implementation chooses to use nn.shape[-1] // head_dimension heads while in other implementations 8 heads are used. Is this a test-and-tell case?

[leondgarse]

head_dimension is the head size mentioned in the paper, scaling value in 32, 64, 128. I think 8 is never mentioned. Besides, I have my CoAtNet0 trained on 224 reached accuracy 81.93, but need another run test, may upload it soon.

[leondgarse]

Uh, sorry, understood now. It's some consideration of the parameters count:

Model	Params in paper	nn.shape[-1] // head_dimension heads	8 heads
CoAtNet0	25M	23.8M	22.6M
CoAtNet1	42M	41.0M	38.8M
CoAtNet2	75M	72.6M	67.3M
CoAtNet3	168M	162.3M	147.1M
CoAtNet4	275M	270.7M	244.3M

These's a gap between my definition and the paper reported, just using num_heads=8 would have a larger gap, but not tested training.

[leondgarse]

This is one of my references while implementing mine: Github comvex/coatnet/model.py, and they are using nn.shape[-1] // head_dimension heads in Github comvex/utils/multihead_attention.py. You may also refer this for pytorch.

[xskxzr]

Could you please show me model.summary() of CoAtNet0? I want to see the number of parameters for each layer but I cannot run your code directly for some reasons.

[leondgarse]

Just run in colab: efficientnetV2_basic_test.ipynb. The last CoAtNet part.

[leondgarse]

split has an param num_or_size_splits specifies the split shapes: query, key, value = tf.split(qkv, [qk_out, qk_out, out_shape], axis=-1), not an even split. This is a common usage handling qk_out != out_shape cases which you can also find in timm.