r""" Deterministic Parameter-Free Projection (DPFP) feature map in `Linear Transformers Are Secretly Fast Weight Programmers `_
| 95 | |
| 96 | |
| 97 | class DPFPFeatureMap(nn.Module): |
| 98 | |
| 99 | r""" |
| 100 | Deterministic Parameter-Free Projection (DPFP) feature map in |
| 101 | `Linear Transformers Are Secretly Fast Weight Programmers <https://arxiv.org/abs/2102.11174>`_ |
| 102 | """ |
| 103 | |
| 104 | def __init__( |
| 105 | self, |
| 106 | head_dim: int, |
| 107 | nu: int = 4 |
| 108 | ) -> DPFPFeatureMap: |
| 109 | super().__init__() |
| 110 | self.nu = nu |
| 111 | |
| 112 | def forward(self, x: torch.Tensor): |
| 113 | x = torch.cat([x.relu(), -x.relu()], dim=-1) |
| 114 | x_rolled = torch.cat([x.roll(shifts=j, dims=-1) for j in range(1, self.nu+1)], dim=-1) |
| 115 | x_repeat = torch.cat([x] * self.nu, dim=-1) |
| 116 | return x_repeat * x_rolled |
| 117 | |
| 118 | |
| 119 | class HadamardFeatureMap(nn.Module): |