| import torch |
| from torch import nn, Tensor |
|
|
|
|
|
|
| class VecDyT(nn.Module): |
| def __init__(self, input_shape): |
| |
| super().__init__() |
| |
| self.alpha = nn.Parameter(torch.randn(input_shape)) |
| |
| def forward(self, x): |
| x = torch.tanh(self.alpha * x) |
| return x |
| |
| |
| class GatingUnit(nn.Module): |
| def __init__(self,dim): |
| |
| super().__init__() |
|
|
| self.proj_1 = nn.Linear(dim,dim,bias=False) |
| self.proj_2 = nn.Linear(dim,dim,bias=False) |
| |
| self.gelu = nn.GELU() |
| |
| |
| def forward(self, x): |
|
|
| u, v = x, x |
| u = self.proj_1(u) |
| u = self.gelu(u) |
| v = self.proj_2(v) |
| g = u * v |
| |
| return g |
|
|
| class TTT(nn.Module): |
| def __init__(self, dim: int): |
| |
| super(TTT, self).__init__() |
| |
| |
| self.mapping = nn.Linear(dim,dim,bias=False) |
| self.State = nn.Linear(dim,dim,bias=False) |
| self.Probe = nn.Linear(dim,dim,bias=False) |
| |
| |
| def forward(self, in_seq: Tensor) -> Tensor: |
|
|
| |
| outs = [] |
| |
| for seq in range(in_seq.size(1)): |
| |
| state = self.State(in_seq[:,seq,:]) |
| train_view = state + torch.randn_like(state) |
| label_view = state |
| loss = nn.functional.mse_loss(self.mapping(train_view), label_view) |
| grads = torch.autograd.grad( |
| loss, self.mapping.parameters(),create_graph=True) |
| with torch.no_grad(): |
| for param, grad in zip(self.mapping.parameters(), grads): |
| |
| param -= 0.01 * grad |
| readout = self.mapping(self.Probe(in_seq[:,seq,:])).detach() |
| outs.append(readout) |
| out = torch.stack(outs, dim=1) |
| |
| return out |
| |
|
|
|
|
| class TensorMapperBlock(nn.Module): |
| def __init__(self, dim, num_patch): |
| |
| super().__init__() |
| |
| self.norm_1 = VecDyT(dim) |
| self.norm_2 = VecDyT(dim) |
| self.memory = TTT(dim) |
| self.feedforward = GatingUnit(dim) |
| |
|
|
| def forward(self, x): |
| |
| |
| |
| residual = x |
| |
| x = self.norm_1(x) |
| |
| x = self.memory(x) |
| |
| x = x + residual |
| |
| residual = x |
| |
| x = self.norm_2(x) |
| |
| x = self.feedforward(x) |
| |
| x = x + residual |
|
|
| return x |
|
|
|
|
| class TensorMapper(nn.Module): |
| def __init__(self, d_model,num_patch, num_layers): |
| super().__init__() |
| |
| self.model = nn.Sequential( |
| *[TensorMapperBlock(d_model,num_patch) for _ in range(num_layers)] |
| ) |
|
|
| def forward(self, x): |
| |
| return self.model(x) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
