|
""" Alibi position bias """ |
|
import math |
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
|
|
def pad_at_dim(t, pad, dim=-1, value=0.0): |
|
dims_from_right = (-dim - 1) if dim < 0 else (t.ndim - dim - 1) |
|
zeros = (0, 0) * dims_from_right |
|
return F.pad(t, (*zeros, *pad), value=value) |
|
|
|
|
|
class AlibiPositionalBias(nn.Module): |
|
def __init__(self, heads, **kwargs): |
|
super().__init__() |
|
self.heads = heads |
|
slopes = torch.Tensor(self._get_slopes(heads)) |
|
slopes = slopes.unsqueeze(1).unsqueeze(1) |
|
self.register_buffer("slopes", slopes, persistent=False) |
|
self.register_buffer("bias", None, persistent=False) |
|
|
|
def get_bias(self, i, j, device): |
|
i_arange = torch.arange(j - i, j, device=device) |
|
j_arange = torch.arange(j, device=device) |
|
bias = -torch.abs( |
|
j_arange.unsqueeze(0).unsqueeze(0) - i_arange.unsqueeze(1).unsqueeze(0) |
|
) |
|
return bias |
|
|
|
@staticmethod |
|
def _get_slopes(heads): |
|
def get_slopes_power_of_2(n): |
|
start = 2 ** (-(2 ** -(math.log2(n) - 3))) |
|
ratio = start |
|
return [start * ratio**i for i in range(n)] |
|
|
|
if math.log2(heads).is_integer(): |
|
return get_slopes_power_of_2(heads) |
|
closest_power_of_2 = 2 ** math.floor(math.log2(heads)) |
|
|
|
return ( |
|
get_slopes_power_of_2(closest_power_of_2) |
|
+ get_slopes_power_of_2(2 * closest_power_of_2)[0::2][ |
|
: heads - closest_power_of_2 |
|
] |
|
) |
|
|
|
def forward(self, qk_dots): |
|
h, i, j, device = *qk_dots.shape[-3:], qk_dots.device |
|
|
|
if (self.bias is not None) and self.bias.shape[-1] >= j: |
|
return qk_dots + self.bias[..., :i, :j] |
|
|
|
bias = self.get_bias(i, j, device) |
|
bias = bias * self.slopes |
|
|
|
num_heads_unalibied = h - bias.shape[0] |
|
bias = pad_at_dim(bias, (0, num_heads_unalibied), dim=0) |
|
self.register_buffer("bias", bias, persistent=False) |
|
|
|
return qk_dots + self.bias |
|
|
