| 204 | |
| 205 | |
| 206 | class RebasedFeatureMap(nn.Module): |
| 207 | |
| 208 | def __init__( |
| 209 | self, |
| 210 | head_dim: int, |
| 211 | use_gamma: Optional[bool] = True, |
| 212 | use_beta: Optional[bool] = True, |
| 213 | normalize: Optional[bool] = True |
| 214 | ) -> RebasedFeatureMap: |
| 215 | super().__init__() |
| 216 | |
| 217 | self.head_dim = head_dim |
| 218 | self.use_gamma = use_gamma |
| 219 | self.use_beta = use_beta |
| 220 | self.normalize = normalize |
| 221 | |
| 222 | self.gamma = None |
| 223 | self.beta = None |
| 224 | if use_gamma: |
| 225 | self.gamma = nn.Parameter(torch.ones(head_dim)) |
| 226 | if use_beta: |
| 227 | self.beta = nn.Parameter(torch.zeros(head_dim)) |
| 228 | |
| 229 | def forward(self, x: torch.Tensor, flatten: Optional[bool] = True): |
| 230 | if self.use_beta and self.use_gamma and self.normalize: |
| 231 | x = layer_norm_fn(x, self.gamma, self.beta) |
| 232 | elif self.normalize: |
| 233 | x = F.layer_norm(x, (self.head_dim,), self.gamma, self.beta) |
| 234 | elif self.use_gamma and self.use_beta: |
| 235 | x = torch.addcmul(self.beta, x, self.gamma) |
| 236 | elif self.use_gamma: |
| 237 | x = x.mul(self.gamma) |
| 238 | else: |
| 239 | raise RuntimeError(f"Not supported combination of `use_gamma`, `use_beta` and `normalize`, " |
| 240 | f"which is currentlt set as (`{self.use_gamma}`, `{self.use_beta}`, `{self.normalize}`)") |
| 241 | if not flatten: |
| 242 | return x |
| 243 | x2_1, x2_2 = flatten_diag_outer_product_off1(x, x) |
| 244 | # rebased use learnable parameters to approximate any quadratic function |
| 245 | return torch.cat([x2_2 * self.head_dim ** -0.5, x2_1 * (2 / self.head_dim) ** 0.5], dim=-1) |
| 246 | |
| 247 | |
| 248 | class ReLUFeatureMap(nn.Module): |