HaloAttention掩码问题

Author: L-wwwwCreated Dec 19, 2024Updated Dec 19, 2024

mask out padding (in the paper, they claim to not need masks, but what about padding?)

    mask = torch.ones(1, 1, h, w, device = device)
    mask = F.unfold(mask, kernel_size = block + (halo * 2), stride = block, padding = halo)
    mask = repeat(mask, '() j i -> (b i h) () j', b = b, h = heads)
    mask = mask.bool()

    max_neg_value = -torch.finfo(sim.dtype).max
    sim.masked_fill_(mask, max_neg_value)

代码似乎将所有非0值置为无穷小,根据注释似乎是想将padding位置 置为无穷小

Source: xmu-xiaoma666/External-Attention-pytorch