feat : add hawk model

model/hawk.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def get_model_device(model):
+    return next(iter(model.parameters())).device
+
+
+class RGLRU(nn.Module):
+    def __init__(self, hidden_size: int, c: float = 8.0):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.c = c
+
+        self.input_gate = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.recurrence_gate = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.a = nn.Parameter(torch.empty(hidden_size))
+
+    def forward(self, x_t: torch.Tensor, state: torch.Tensor) -> torch.Tensor:
+        batch_size, hidden_size = x_t.shape
+        assert hidden_size == self.hidden_size
+        assert state.shape[0] == batch_size
+
+        i_t = torch.sigmoid(self.input_gate(x_t))
+        r_t = torch.sigmoid(self.recurrence_gate(x_t))
+
+        # Compute recurrence
+        a_t = self.a ** (self.c * r_t)
+        multiplier = torch.sqrt(1 - a_t**2)
+        new_state = (state * a_t) + (multiplier * (i_t * x_t))
+
+        return new_state
+
+    def init_state(self, batch_size: int, device: torch.device | None = None):
+        if device is None:
+            device = get_model_device(self)
+        return torch.zeros(batch_size, self.hidden_size, device=device)
+
+
+class CausalConv1d(nn.Module):
+
+    def __init__(self, hidden_size, kernel_size):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.kernel_size = kernel_size
+        self.conv = nn.Conv1d(
+            hidden_size, hidden_size, kernel_size, groups=hidden_size, bias=True
+        )
+
+    def init_state(self, batch_size: int, device: torch.device | None = None):
+        if device is None:
+            device = get_model_device(self)
+        return torch.zeros(
+            batch_size, self.hidden_size, self.kernel_size - 1, device=device
+        )
+
+    def forward(self, x: torch.Tensor, state: torch.Tensor):
+        x_with_state = torch.concat([state, x[:, :, None]], dim=-1)
+        out = self.conv(x_with_state)
+        new_state = x_with_state[:, :, 1:]
+        return out.squeeze(-1), new_state
+
+
+class Hawk(nn.Module):
+    def __init__(self, hidden_size: int, conv_kernel_size: int = 4):
+        super().__init__()
+
+        self.conv_kernel_size = conv_kernel_size
+        self.hidden_size = hidden_size
+
+        self.gate_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.recurrent_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.conv = CausalConv1d(hidden_size, conv_kernel_size)
+        self.rglru = RGLRU(hidden_size)
+        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+
+    def forward(
+        self, x: torch.Tensor, state: tuple[torch.Tensor, torch.Tensor]
+    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        conv_state, rglru_state = state
+
+        batch_size, hidden_size = x.shape
+        
+        assert batch_size == conv_state.shape[0] == rglru_state.shape[0]
+        assert self.hidden_size == hidden_size == rglru_state.shape[1]
+
+        gate = F.gelu(self.gate_proj(x))
+        x = self.recurrent_proj(x)
+
+        x, new_conv_state = self.conv(x, conv_state)
+        new_rglru_state = self.rglru(x, rglru_state)
+
+        gated = gate * new_rglru_state
+        out = self.out_proj(gated)
+
+        new_state = [new_conv_state, new_rglru_state]
+        return out, new_state
+
+    def init_state(
+        self, batch_size: int, device: torch.device | None = None
+    ) -> list[torch.Tensor]:
+        return [
+            self.conv.init_state(batch_size, device),
+            self.rglru.init_state(batch_size, device),
+        ]