| 188 | |
| 189 | |
| 190 | class TaylorFeatureMap(nn.Module): |
| 191 | def __init__( |
| 192 | self, |
| 193 | head_dim: int |
| 194 | ) -> TaylorFeatureMap: |
| 195 | super().__init__() |
| 196 | self.head_dim = head_dim |
| 197 | self.r2 = math.sqrt(2) |
| 198 | self.rd = math.sqrt(self.head_dim) |
| 199 | self.rrd = math.sqrt(self.rd) |
| 200 | |
| 201 | def forward(self, x: torch.Tensor): |
| 202 | x2_1, x2_2 = flatten_diag_outer_product_off1(x, x) |
| 203 | return torch.cat([torch.ones_like(x[..., 0:1]), x / self.rrd, x2_2 / (self.rd * self.r2), x2_1 / self.rd], dim=-1) |
| 204 | |
| 205 | |
| 206 | class RebasedFeatureMap(nn.Module): |