r""" Hedgehog feature map as introduced in `The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry `_
| 36 | |
| 37 | |
| 38 | class HedgehogFeatureMap(nn.Module): |
| 39 | |
| 40 | r""" |
| 41 | Hedgehog feature map as introduced in |
| 42 | `The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry <https://arxiv.org/abs/2402.04347>`_ |
| 43 | """ |
| 44 | |
| 45 | def __init__( |
| 46 | self, |
| 47 | head_dim: int |
| 48 | ) -> HedgehogFeatureMap: |
| 49 | super().__init__() |
| 50 | # Trainable map |
| 51 | self.layer = nn.Linear(head_dim, head_dim) |
| 52 | self.init_weights_() |
| 53 | |
| 54 | def init_weights_(self): |
| 55 | """Initialize trainable map as identity""" |
| 56 | with torch.no_grad(): |
| 57 | identity = torch.eye(*self.layer.weight.shape[-2:], dtype=torch.float) |
| 58 | self.layer.weight.copy_(identity.to(self.layer.weight)) |
| 59 | nn.init.zeros_(self.layer.bias) |
| 60 | |
| 61 | def forward(self, x: torch.Tensor): |
| 62 | x = self.layer(x) # shape b, h, l, d |
| 63 | return torch.cat([2*x, -2*x], dim=-1).softmax(-1) |
| 64 | |
| 65 | |
| 66 | class T2RFeatureMap(nn.Module): |