same for first

Author: davidweichiang

first.py

@@ -7,6 +7,17 @@
 
 log_sigmoid = torch.nn.LogSigmoid()
 
+class PositionEncoding(torch.nn.Module):
+    def __init__(self, size):
+        super().__init__()
+        self.size = size
+
+    def forward(self, n):
+        zero = torch.zeros(n)
+        pos = torch.arange(0, n).to(torch.float)
+        pe = torch.stack([pos == 1] + [zero]*(self.size-1), dim=1)
+        return pe
+
 ap = argparse.ArgumentParser()
 ap.add_argument('--train_length', type=int, default=50)
 ap.add_argument('--test_length', type=int, default=1000)
@@ -17,30 +28,24 @@
 
 alphabet = ["0", "1", "$"]
 alphabet_index = {a:i for i,a in enumerate(alphabet)}
-max_pos = 10000
 size = 16
 
 class Model(torch.nn.Module):
     def __init__(self, alphabet_size, size):
         super().__init__()
 
         self.word_embedding = torch.nn.Embedding(num_embeddings=alphabet_size, embedding_dim=size)
-        self.pos_embedding = torch.stack([
-            torch.arange(0, max_pos, dtype=torch.float) == 0,
-            torch.arange(0, max_pos, dtype=torch.float) == 1,
-            torch.arange(0, max_pos, dtype=torch.float) >= 2,
-        ], dim=1).to(torch.float)
-        self.pos_adapter = torch.nn.Linear(self.pos_embedding.size()[1], size)
+        self.pos_encoding = PositionEncoding(size)
 
-        encoder_layer = encoder.PostnormTransformerEncoderLayer(d_model=size, nhead=1, dim_feedforward=size*4, dropout=0.)
+        encoder_layer = encoder.TransformerEncoderLayer(d_model=size, nhead=1, dim_feedforward=size*4, dropout=0.)
         #encoder_layer = encoder.ScaledTransformerEncoderLayer(d_model=size, nhead=1, dim_feedforward=size*4, dropout=0.)
         #encoder_layer.norm1.eps = encoder_layer.norm2.eps = 0.
         self.encoder = torch.nn.TransformerEncoder(encoder_layer, num_layers=2)
 
         self.output_layer = torch.nn.Linear(size, 1)
 
     def forward(self, w):
-        x = self.word_embedding(w) + self.pos_adapter(self.pos_embedding[:len(w)])
+        x = self.word_embedding(w) + self.pos_encoding(len(w))
         y = self.encoder(x.unsqueeze(1)).squeeze(1)
         y = y[0]
         z = self.output_layer(y)

first_exact.py

@@ -5,8 +5,6 @@
 import sys
 import argparse
 
-log_sigmoid = torch.nn.LogSigmoid()
-
 ap = argparse.ArgumentParser()
 ap.add_argument('--length', type=int, default=100)
 ap.add_argument('--steps', type=int, default=100)
@@ -17,6 +15,21 @@
 alphabet_index = {a:i for i,a in enumerate(alphabet)}
 max_pos = 10000
 
+log_sigmoid = torch.nn.LogSigmoid()
+
+class PositionEncoding(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, n):
+        zero = torch.zeros(n)
+        pos = torch.arange(0, n).to(torch.float)
+        pe = torch.stack([zero]*3 +
+                         [pos == 1] +
+                         [zero]*2,
+                         dim=1)
+        return pe
+
 class FirstLayer(torch.nn.TransformerEncoderLayer):
     def __init__(self):
         super().__init__(6, 1, 1, dropout=0.)
@@ -96,20 +109,15 @@ def __init__(self):
         super().__init__()
 
         self.word_embedding = torch.eye(3, 6)
-        self.pos_embedding = torch.stack(
-            [torch.zeros(max_pos)]*3 +
-            [torch.arange(0, max_pos, dtype=torch.float) == 1] +
-            [torch.zeros(max_pos)]*2,
-            dim=1)
-
+        self.pos_encoding = PositionEncoding()
         self.transformer_encoder = MyTransformerEncoder()
         self.output_layer = torch.nn.Linear(6, 1)
         self.output_layer.weight = torch.nn.Parameter(torch.tensor(
             [[0,0,0,0,0,1]], dtype=torch.float))
         self.output_layer.bias = torch.nn.Parameter(torch.tensor([0.]))
 
     def forward(self, w):
-        x = self.word_embedding[w] + self.pos_embedding[:len(w)]
+        x = self.word_embedding[w] + self.pos_encoding(len(w))
         y = self.transformer_encoder(x.unsqueeze(1)).squeeze(1)
         z = self.output_layer(y[0])
         return z

first_exact_layernorm.py

@@ -5,8 +5,6 @@
 import sys
 import argparse
 
-log_sigmoid = torch.nn.LogSigmoid()
-
 ap = argparse.ArgumentParser()
 ap.add_argument('--length', type=int, default=100)
 ap.add_argument('--steps', type=int, default=100)
@@ -19,6 +17,21 @@
 alphabet_index = {a:i for i,a in enumerate(alphabet)}
 max_pos = 10000
 
+log_sigmoid = torch.nn.LogSigmoid()
+
+class PositionEncoding(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, n):
+        zero = torch.zeros(n)
+        pos = torch.arange(0, n).to(torch.float)
+        pe = torch.stack([zero]*3 +
+                         [pos == 1] +
+                         [zero]*2,
+                         dim=1)
+        return pe
+
 class FirstLayer(torch.nn.TransformerEncoderLayer):
     def __init__(self):
         super().__init__(12, 1, 1, dropout=0.)
@@ -110,20 +123,15 @@ def __init__(self):
         super().__init__()
 
         self.word_embedding = torch.eye(3, 6)
-        self.pos_embedding = torch.stack(
-            [torch.zeros(max_pos)]*3 +
-            [torch.arange(0, max_pos, dtype=torch.float) == 1] +
-            [torch.zeros(max_pos)]*2,
-            dim=1)
-
+        self.pos_encoding = PositionEncoding()
         self.transformer_encoder = MyTransformerEncoder()
         self.output_layer = torch.nn.Linear(12, 1)
         self.output_layer.weight = torch.nn.Parameter(torch.tensor(
             [[0,0,0,0,0,1,0,0,0,0,0,0]], dtype=torch.float))
         self.output_layer.bias = torch.nn.Parameter(torch.tensor([0.]))
 
     def forward(self, w):
-        x = self.word_embedding[w] + self.pos_embedding[:len(w)]
+        x = self.word_embedding[w] + self.pos_encoding(len(w))
         x = torch.cat([x, -x], dim=-1)
         y = self.transformer_encoder(x.unsqueeze(1)).squeeze(1)
         z = self.output_layer(y[0])

parity_exact_layernorm.py

@@ -179,9 +179,8 @@ def __init__(self):
         self.output_layer.bias = torch.nn.Parameter(torch.tensor([0.]))
 
     def forward(self, w):
-        x = torch.cat([self.word_embedding[w] + self.pos_encoding(len(w)),
-                       -(self.word_embedding[w] + self.pos_encoding(len(w)))],
-                      dim=1)
+        x = self.word_embedding[w] + self.pos_encoding(len(w))
+        x = torch.cat([x, -x], dim=-1)
         y = self.transformer_encoder(x.unsqueeze(1)).squeeze(1)
         z = self.output_layer(y[-1])
         return z