Update.
[mygptrnn.git] / mygpt.py
index 7c9991f..492a9bb 100755 (executable)
--- a/mygpt.py
+++ b/mygpt.py
@@ -126,7 +126,6 @@ class AddPositionalEncoding(nn.Module):
 
 import pscan
 
-
 # X is /.../xTxD   A is /.../xT   Y_init is /.../xD
 
 
@@ -147,6 +146,18 @@ def pscan_dim(A, X, Y_init, dim=-2):
     return Y
 
 
+def pscan_rgrad(grad_Y, A, X, Y_init, dim=-2, eps=1e-2):
+    with torch.no_grad():
+        s_A, s_X = 0, 0
+        for t in range(X.size(dim) - 1, 0, -1):
+            delta = (grad_Y[t] - s_A) / A[t].grad
+            s_A += A[t].grad * delta
+            A[t].grad = delta
+            delta = (grad_Y[t] - s_X) / X[t].grad
+            s_X += X[t].grad * delta
+            X[t].grad = delta
+
+
 def pscan_shape(A, X, Y_init):
     s = X.size()
     A = A.reshape(-1, s[-2])
@@ -491,16 +502,29 @@ class Caterpillar(nn.Module):
         self.caterpillar_height = caterpillar_height
         self.attention_dropout = attention_dropout
 
-        self.proba_gate_dropout = 0.0
+        ######################################################################
+        # sup_args
 
-        default_b_G = kwargs.get("default_b_G")
-        if default_b_G is None:
-            default_b_G = -math.log(caterpillar_height - 1)
+        x = kwargs.get("gate_dropout")
+        if x is None:
+            self.proba_gate_dropout = 0.0
+        else:
+            self.proba_gate_dropout = float(x)
+
+        logger(f"self.proba_gate_dropout {self.proba_gate_dropout}")
 
-        logger(f"default_b_G {default_b_G}")
+        x = kwargs.get("default_bg")
+        if x is None:
+            default_bg = -math.log(caterpillar_height - 1)
+        else:
+            default_bg = float(x)
+
+        logger(f"default_bg {default_bg}")
+
+        ######################################################################
 
         self.w_G = randw(nb_heads, caterpillar_height, dim_model)
-        self.b_G = nn.Parameter(torch.full((nb_heads, caterpillar_height), default_b_G))
+        self.b_G = nn.Parameter(torch.full((nb_heads, caterpillar_height), default_bg))
 
         self.w_K = randw(nb_heads, dim_qk, dim_model)
         self.w_V = randw(nb_heads, dim_v, dim_model)
@@ -552,8 +576,8 @@ class Caterpillar(nn.Module):
             self.rec_K = X.new_zeros(N, R, T, DK)
             # We start the recurrent sequences with optimizable
             # initial values. No idea if it helps.
-            self.rec_V[:, :, t0 - L : t0] = self.init_V_rec[None, :, :, :]
-            self.rec_K[:, :, t0 - L : t0] = self.init_K_rec[None, :, :, :]
+            self.rec_V[:, :, t0 - L : t0, :] = self.init_V_rec[None, :, :, :]
+            self.rec_K[:, :, t0 - L : t0, :] = self.init_K_rec[None, :, :, :]
 
             self.cache_Y = X.new_zeros(N, T, DM)
 
@@ -579,84 +603,66 @@ class Caterpillar(nn.Module):
         G = G / G.sum(1, keepdim=True).clamp(min=1)
 
         ######################################################################
-        # Roll the gating indexes
 
-        # warnings.warn("rotating barrel", RuntimeWarning)
+        def recurrence(G, V, K):
+            # We prepare the arguments for the parallel scan
 
-        # r_barrel = torch.arange(R, device=G.device)[None, None, :, None]
-        # t_barrel = torch.arange(t1 - t0, device=G.device)[None, None, None, :]
-        # r_barrel = (r_barrel + (t_barrel + t0) // L) % R
-        # G = G.gather(dim=2, index=r_barrel.expand_as(G))
+            A = 1 - G.sum(1)
 
-        ######################################################################
-        # The "flashbacks"
+            gated_V = torch.einsum("nhrt,nhtd->nrtd", G, V)
+            gated_K = torch.einsum("nhrt,nhtd->nrtd", G, K)
 
-        if self.training and self.proba_gate_dropout > 0.0:
-            # This is a better implementation of "flashbacks".
+            # We start from cached values, which matters in inference
 
-            # G is NxHxExT where e is the caterpillar's row.
+            init_rec_V = self.rec_V[:, :, t0 - L : t0]
+            init_rec_K = self.rec_K[:, :, t0 - L : t0]
 
-            warnings.warn("gate dropout", RuntimeWarning)
-            epsilon = 0.5
+            # Associative scan
 
-            dropout_head = (
-                (torch.rand(N, H, 1, t1 - t0, device=G.device).sort(dim=3).indices == 0)
-                .expand_as(G)
-                .float()
-            )
-
-            dropout_tail = dropout_head.cumsum(dim=3) - dropout_head
-
-            dropout_active = (
-                torch.rand(N, 1, 1, 1, device=G.device) < self.proba_gate_dropout
-            ).long()
+            # Here there is a trick: Since the stack at position t is
+            # computed by updating that at position t-L, the parallel
+            # scan operates with a period of L. To do so we split the
+            # sequence indexing in two axes, the second of size L, and
+            # run the parallel scan using the first as the sequence index.
 
-            dropout_head *= dropout_active
-            dropout_tail *= dropout_active
+            A = A.unflatten(2, (-1, L))
+            gated_V = gated_V.unflatten(2, (-1, L))
+            gated_K = gated_K.unflatten(2, (-1, L))
 
-            G = (
-                G
-                + dropout_head * (1 - epsilon - G.detach())
-                - dropout_tail * G.detach()
-            )
-
-        ######################################################################
+            next_V = pscan_dim(A, gated_V, init_rec_V, dim=2)
+            next_K = pscan_dim(A, gated_K, init_rec_K, dim=2)
 
-        # We prepare the arguments for the parallel scan
+            next_V = next_V.flatten(2, 3)
+            next_K = next_K.flatten(2, 3)
 
-        A = 1 - G.sum(1)
+            return next_V, next_K
 
-        # warnings.warn("harmonic recurrence", RuntimeWarning)
-        # har = torch.arange(t0, t1, device = G.device).float() + 1
-        # A = har / (har + 1)
-        # G = G / har
+        #################################################################
 
-        gated_V = torch.einsum("nhrt,nhtd->nrtd", G, V)
-        gated_K = torch.einsum("nhrt,nhtd->nrtd", G, K)
+        next_V, next_K = recurrence(G, V, K)
 
-        # We start from cached values, which matters in inference
+        if self.training and self.proba_gate_dropout > 0.0:
+            # G is NxHxRxT where r is the caterpillar's row.
 
-        init_rec_V = self.rec_V[:, :, t0 - L : t0]
-        init_rec_K = self.rec_K[:, :, t0 - L : t0]
+            warnings.warn("gate dropout", RuntimeWarning)
 
-        #################################################################
-        # Associative scan
+            kill = (
+                torch.rand(G.size(), device=G.device) <= self.proba_gate_dropout
+            ).float()
 
-        # Here there is a trick: Since the stack at position t is
-        # computed by updating that at position t-L, the parallel
-        # scan operates with a period of L. To do so we split the
-        # sequence indexing in two axes, the second of size L, and
-        # run the parallel scan using the first as the sequence index.
+            mask = 1 - kill
 
-        A = A.unflatten(2, (-1, L))
-        gated_V = gated_V.unflatten(2, (-1, L))
-        gated_K = gated_K.unflatten(2, (-1, L))
+            masked_next_V, masked_next_K = recurrence(G * mask, V, K)
 
-        next_V = pscan_dim(A, gated_V, init_rec_V, dim=2)
-        next_K = pscan_dim(A, gated_K, init_rec_K, dim=2)
+            next_V = next_V.detach() + (masked_next_V - masked_next_V.detach()) / (
+                1 - self.proba_gate_dropout
+            )
+            next_K = next_K.detach() + (masked_next_K - masked_next_K.detach()) / (
+                1 - self.proba_gate_dropout
+            )
 
-        self.rec_V[:, :, t0:t1] = next_V.flatten(2, 3)
-        self.rec_K[:, :, t0:t1] = next_K.flatten(2, 3)
+        self.rec_V[:, :, t0:t1] = next_V
+        self.rec_K[:, :, t0:t1] = next_K
 
         ######################################################################
         # compute the readout