Adding nonlinear projection in self attention

mingpt/model.py

@@ -83,7 +83,12 @@ class CausalSelfAttention(nn.Module):
         # key, query, value projections for all heads, but in a batch
         # self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
         # output projection
-        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
+        #self.c_proj = nn.Linear(config.n_embd, config.n_embd)
+        self.c_proj = nn.Sequential(
+            nn.Linear(config.n_embd, config.n_embd),
+            nn.ReLU(),
+            nn.Linear(config.n_embd, config.n_embd)
+        )
         # regularization
         self.attn_dropout = nn.Dropout(config.attn_pdrop)
         self.resid_dropout = nn.Dropout(config.resid_pdrop)
@@ -102,8 +107,8 @@ class CausalSelfAttention(nn.Module):
             y[mask_tokens] = 0
 
         # output projection
-        # y = self.resid_dropout(self.c_proj(y))
-        y = self.resid_dropout(y)
+        y = self.resid_dropout(self.c_proj(y))
+        # y = self.resid_dropout(y)
         return y