PosEmbed is different for each block - shouldn't it be the same!?

[jxtps]

Reading https://github.com/google-research/nested-transformer/blob/main/models/nest_net.py#L89 and https://github.com/google-research/nested-transformer/blob/main/libml/self_attention.py#L225 it's clear that the PositionEmbedding is done over three not two axes.

This is further corroborated by the pytorch clone: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/nest.py#L204

So it's taken over (blocks, seqlen, d), not (1, seqlen, d), which means that each block (= part of the image) gets its own set of pos embeddings.

Is this intentional? It seems to run counter to the spirit of the paper?!

With (1, seqlen, d) there would be more of an inductive bias, right? With (blocks, seqlen, d) the earlier stages can "cheat" and do different things in different blocks (= parts of) the image?

And more importantly, the separate blocks don't boost each other's modeling of the pos embedding - wouldn't you expect even faster convergence (arguably the paper's main achievement) with (1, seqlen, d)?

Thanks!

(This only matters on the first two levels in the hierarchy where #blocks > 1)