Stanford CS 312 | Deep Learning Alchemy

Example Assignment Problem: Start from the baseline depth 8 transformer that we pre-train on 614M tokens. This transformer (implementation provided below) features multiple components such as prenorm, attention, MLP, residuals, and positional embeddings. In this problem, we want to understand how much each component improves loss. Explore the following directions.

Perform leave-one-out ablations on the role of each component in our current model.
...
...
Understand the role of the residual connection. Can you come up with different schemes that improve the loss benefit of the residual connection?

Relevant code snippet

class Block(nn.Module):
    def __init__(self, width, ...):
        super().__init__()
        self.attn_norm = RMSNorm(width)
        self.attn = Attention(width, ...)
        self.mlp_norm = RMSNorm(width)
        self.mlp = MLP(width, ...)

    def update(self, x):
        attn = self.attn(self.attn_norm(x))
        mlp = self.mlp(self.mlp_norm(x + attn))
        return attn + mlp

    def forward(self, x):
        return x + self.update(x)


class Transformer(nn.Module):
    def __init__(self, vocab_size, depth=8, ...):
        super().__init__()
        self.token_embed = nn.Embedding(vocab_size, ...)
        self.position_embed = nn.Embedding(...)
        self.blocks = nn.ModuleList(
            [Block(...) for _ in range(depth)]
        )
        self.final_norm = RMSNorm(...)
        self.lm_head = nn.Linear(..., vocab_size)

    def forward(self, tokens):
        positions = torch.arange(
            tokens.size(1), device=tokens.device
        )
        x = self.token_embed(tokens)
        x = x + self.position_embed(positions)
        for block in self.blocks:
            x = block(x)
        return self.lm_head(self.final_norm(x))

Example Quiz Question (Difficulty Easy): We train each of the following depth-8 transformer variants using the same training recipe and hyperparameters. Rank the variants from lowest to highest final validation loss. Use <, >, or = between the letters; use = when the losses are within 0.01.

Baseline
No prenorm
No residual
No MLP
No attention
No positional embeddings

No prenorm diff

@@ -10,4 +10,4 @@ class Block.update
 def update(self, x):
-    attn = self.attn(self.attn_norm(x))
+    attn = self.attn(x)
-    mlp = self.mlp(self.mlp_norm(x + attn))
+    mlp = self.mlp(x + attn)
     return attn + mlp

No residual diff

@@ -16,2 +16,2 @@ class Block.forward
 def forward(self, x):
-    return x + self.update(x)
+    return self.update(x)

No MLP diff

@@ -10,4 +10,3 @@ class Block.update
 def update(self, x):
     attn = self.attn(self.attn_norm(x))
-    mlp = self.mlp(self.mlp_norm(x + attn))
-    return attn + mlp
+    return attn

No attention diff

@@ -10,4 +10,3 @@ class Block.update
 def update(self, x):
-    attn = self.attn(self.attn_norm(x))
-    mlp = self.mlp(self.mlp_norm(x + attn))
-    return attn + mlp
+    mlp = self.mlp(self.mlp_norm(x))
+    return mlp

No positional embeddings diff

@@ -22,2 +22,1 @@ Transformer.__init__
         self.token_embed = nn.Embedding(vocab_size, ...)
-        self.position_embed = nn.Embedding(...)
@@ -30,8 +29,4 @@ Transformer.forward
     def forward(self, tokens):
-        positions = torch.arange(
-            tokens.size(1), device=tokens.device
-        )
         x = self.token_embed(tokens)
-        x = x + self.position_embed(positions)
         for block in self.blocks:
             x = block(x)

Example Quiz Question (Difficulty Hard): We train the default 8 layer transformer as well as a variant where instead of having the residual connection come from the previous layer, we take 0.5 contribution from the previous layer and 0.5 contribution from two layers before (except for the first layer). For each method, we take the best loss over six learning rates {10⁻⁴, 3 × 10⁻⁴, 10⁻³, 3 × 10⁻³, 10⁻², 3 × 10⁻²}. What is your prediction of the loss difference between each method (L_{mixed residual} − L_default)?

Core diff

@@ -34,3 +34,9 @@ Transformer.forward
 x_two_back = None
 for i, block in enumerate(self.blocks):
-    x = block(x)
+    update = block.update(x)
+    residual = (
+        x if i == 0
+        else 0.5 * x + 0.5 * x_two_back
+    )
+    x_two_back, x = x, residual + update

Note: both of these questions are hard, do not get discouraged if you do not have intuition for them without running experiments from the assignment!