Confused about the shape of relative position encoding

[Jensen-Su]

The following code generates a position embedding of shape (C, K, K), where C=self.group_planes*2, K=self.kernel_size:

axial-deeplab/lib/models/axialnet.py

Line 64 in fe1d052

all_embeddings = torch.index_select(self.relative, 1, self.flatten_index).view(self.group_planes * 2, self.kernel_size, self.kernel_size)

It seems that each position in the (K, K) window owns a position encoding. But for axial-attention applied along w-axis, shouldn't the shape be (C, W), meaning that all rows share a same position encoding ?

[csrhddlam]

Does (C, W) mean global position encoding?
Note that our positional encoding is relative and shared across the other axis. For example, in w-axis attention, each pixel corresponds to W other pixels and W relative positions, so there are (W, W) relative positional encodings in total.

[Jensen-Su]

Does (C, W) mean global position encoding?
Note that our positional encoding is relative and shared across the other axis. For example, in w-axis attention, each pixel corresponds to W other pixels and W relative positions, so there are (W, W) relative positional encodings in total.

Thanks for your helpful reply. I did mean global position encoding by (C, W).
I am also confused about the following line:

axial-deeplab/lib/models/axialnet.py

Line 44 in fe1d052

relative_index = key_index - query_index + kernel_size - 1

The confusion is that, since all the position encodings are initialized randomly, I expected that whatever orders we index the relative encoding should result in similar results. So maybe we can index it with a simpler way. But clearly you don't think so by using this relative_index. What do I miss?

[phj128]

They are randomly initialized, but for different position they have different relative positional encoding while the same relative distance ones should share the weights.

[mcahny]

Note that our positional encoding is relative and shared across the other axis. For example, in w-axis attention, each pixel corresponds to W other pixels and W relative positions, so there are (W, W) relative positional encodings in total.

If the span-size K is smaller than the width W, then do we have the size of (C,W,K) for the relative position encoding matrix?
So that it's einsumed with the query like Q (H,(W,C)) * r^q ((W,C),K) -> A (H,W,K)? (A: attention matrix)