Update.
[mygptrnn.git] / mygpt.py
index 3a48cdb..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])
@@ -565,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)
 
@@ -586,81 +597,72 @@ class Caterpillar(nn.Module):
             torch.einsum("ntc,hrc->nhrt", X, self.w_G) + self.b_G[None, :, :, None]
         ).sigmoid()
 
-        # warnings.warn("softmax gating", RuntimeWarning)
+        # Clip the gating to avoid values greater than 1 when several
+        # heads hit the same row
 
-        # G = (
-        # torch.einsum("ntc,hrc->nhrt", X, self.w_G) + self.b_G[None, :, :, None]
-        # ).softmax(dim=2)
+        G = G / G.sum(1, keepdim=True).clamp(min=1)
 
         ######################################################################
-        # The "flashbacks"
 
-        if self.training and self.proba_gate_dropout > 0.0:
-            # This is a better implementation of "flashbacks".
+        def recurrence(G, V, K):
+            # We prepare the arguments for the parallel scan
 
-            # G is NxHxExT where e is the caterpillar's row.
+            A = 1 - G.sum(1)
 
-            warnings.warn("gate dropout", RuntimeWarning)
+            gated_V = torch.einsum("nhrt,nhtd->nrtd", G, V)
+            gated_K = torch.einsum("nhrt,nhtd->nrtd", G, K)
 
-            kill = (
-                torch.rand(G.size(), device=G.device) <= self.proba_gate_dropout
-            ).float()
+            # We start from cached values, which matters in inference
 
-            alpha = G / (1 - self.proba_gate_dropout)
+            init_rec_V = self.rec_V[:, :, t0 - L : t0]
+            init_rec_K = self.rec_K[:, :, t0 - L : t0]
 
-            G = alpha * (1 - kill)
+            # Associative scan
 
-        ######################################################################
-        # Clip the gating to avoid values greater than 1 when several
-        # heads hit the same row
+            # 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.
 
-        G = G / G.sum(1, keepdim=True).clamp(min=1)
-
-        ######################################################################
-        # Roll the gating indexes
-
-        # warnings.warn("rotating barrel", RuntimeWarning)
+            A = A.unflatten(2, (-1, L))
+            gated_V = gated_V.unflatten(2, (-1, L))
+            gated_K = gated_K.unflatten(2, (-1, L))
 
-        # 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))
+            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