Agent Conversation: gemini-3.1-pro-preview
Model: gemini-3.1-pro-preview
Total steps: 8
Messages: 10
UserStep 0
Task: llm-pretrain-attention
LLM Pretraining: Attention Mechanism Optimization
Research Question
Design an improved self-attention mechanism for GPT-2 language model pretraining. Your modifications should reduce validation loss compared to the standard multi-head attention with learned absolute position embeddings.
What You Can Modify
The CausalSelfAttention class (lines 34-70 in custom_pretrain.py), including:
- Position encoding scheme (the default uses learned absolute position embeddings via
wpe) - Query/Key/Value computation and projection
- Attention score computation and masking
- Any attention-related hyperparameters
Note: If your attention mechanism implements its own position encoding (replacing the learned wpe), set self.use_pos_emb = False in __init__ — the model will then skip adding position embeddings in the forward pass.
Evaluation
- Metric: Validation loss (cross-entropy, lower is better)
- Model sizes: GPT-2 124M (12L/12H/768D) and GPT-2 1.5B (48L/25H/1600D, 4-GPU DDP)
- Dataset: FineWeb 10B (GPT-2 tokenizer)
- Training: 5000 iterations, batch_size=12, block_size=1024, grad_accum=5
nanoGPT/custom_pretrain.py [EDITABLE — lines 33–70, lines 245–247 only]
1: """Custom GPT-2 Pretraining Script
2: Based on Andrej Karpathy's nanoGPT, evaluated on FineWeb dataset.
3: """
4:
5: import math
6: import inspect
7: import os
8: import time
9: from contextlib import nullcontext
10: from dataclasses import dataclass
11:
12: import numpy as np
13: import torch
14: import torch.nn as nn
15: from torch.nn import functional as F
16:
17: # ============================================================================
18: # Model Components
19: # ============================================================================
20:
21: # ── Normalization ──────────────────────────────────────────────────────────
22: class LayerNorm(nn.Module):
23: """LayerNorm but with an optional bias."""
24: def __init__(self, ndim, bias):
25: super().__init__()
26: self.weight = nn.Parameter(torch.ones(ndim))
27: self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
28:
29: def forward(self, input):
30: return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
31:
32: # ── Self-Attention ─────────────────────────────────────────────────────────
33: class CausalSelfAttention(nn.Module):
34: def __init__(self, config):
35: super().__init__()
36: assert config.n_embd % config.n_head == 0
37: self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
38: self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
39: self.attn_dropout = nn.Dropout(config.dropout)
40: self.resid_dropout = nAssistantStep 1edit
Tool Input
1{
2 "op": "replace",
3 "filename": "custom_pretrain.py",
4 "start_line": 33,
5 "end_line": 70,
6 "after_line": 1,
7 "content": "class CausalSelfAttention(nn.Module):\n def __init__(self, config):\n super().__init__()\n assert config.n_embd % config.n_head == 0\n self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)\n self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)\n self.attn_dropout = nn.Dropout(config.dropout)\n self.resid_dropout = nn.Dropout(config.dropout)\n self.n_head = config.n_head\n self.n_embd = config.n_embd\n self.head_dim = config.n_embd // config.n_head\n self.dropout = config.dropout\n self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')\n if not self.flash:\n self.register_buffer(\"bias\", torch.tril(torch.ones(config.block_size, config.block_size))\n .view(1, 1, config.block_size, config.block_size))\n \n self.use_pos_emb = False\n \n # RoPE frequencies (using base 100000 for better representation)\n inv_freq = 1.0 / (100000 ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))\n self.register_buffer(\"inv_freq\", inv_freq)\n \n # Token Shift mixing parameters for Q, K, V\n self.time_mix = nn.Parameter(torch.empty(1, 1, 3 * config.n_embd))\n nn.init.uniform_(self.time_mix, 0.2, 0.8)\n\n # Learnable attention temperature per head\n self.attn_scale = nn.Parameter(torch.ones(1, self.n_head, 1, 1))\n\n def _apply_rope(self, x, seq_len):\n t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)\n freqs = torch.outer(t, self.inv_freq)\n cos = freqs.cos().unsqueeze(0).unsqueeze(0)\n sin = freqs.sin().unsqueeze(0).unsqueeze(0)\n d = x.shape[-1] // 2\n x1, x2 = x[..., :d], x[..., d:]\n y1 = x1 * cos - x2 * sin\n y2 = x1 * sin + x2 * cos\n return torch.cat([y1, y2], dim=-1).type_as(x)\n\n def forward(self, x):\n B, T, C = x.size()\n \n qkv = self.c_attn(x)\n \n # Token Shift (Time Mixing)\n qkv_shifted = torch.cat([\n torch.zeros(B, 1, 3 * C, device=qkv.device, dtype=qkv.dtype), \n qkv[:, :-1, :]\n ], dim=1)\n qkv = qkv * self.time_mix + qkv_shifted * (1.0 - self.time_mix)\n \n q, k, v = qkv.split(self.n_embd, dim=2)\n \n k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)\n q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)\n v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)\n \n # QK-Norm without learned weights to preserve RoPE geometry\n q = F.rms_norm(q, (self.head_dim,))\n k = F.rms_norm(k, (self.head_dim,))\n \n # RoPE\n q = self._apply_rope(q, T)\n k = self._apply_rope(k, T)\n \n # Apply learnable temperature per head to Query\n q = q * self.attn_scale\n\n if self.flash:\n y = torch.nn.functional.scaled_dot_product_attention(\n q, k, v, attn_mask=None,\n dropout_p=self.dropout if self.training else 0, is_causal=True)\n else:\n att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))\n att = F.softmax(att, dim=-1)\n att = self.attn_dropout(att)\n y = att @ v\n \n y = y.transpose(1, 2).contiguous().view(B, T, C)\n y = self.resid_dropout(self.c_proj(y))\n return y"
8}Tool ResultStep 1
ERROR: Package 'custom_pretrain.py' is not in allowed packages
AssistantStep 2edit
Tool Input
1{
2 "op": "replace",
3 "after_line": 1,
4 "end_line": 70,
5 "start_line": 33,
6 "filename": "nanoGPT/custom_pretrain.py",
7 "content": "class CausalSelfAttention(nn.Module):\n def __init__(self, config):\n super().__init__()\n assert config.n_embd % config.n_head == 0\n self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)\n self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)\n self.attn_dropout = nn.Dropout(config.dropout)\n self.resid_dropout = nn.Dropout(config.dropout)\n self.n_head = config.n_head\n self.n_embd = config.n_embd\n self.head_dim = config.n_embd // config.n_head\n self.dropout = config.dropout\n self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')\n if not self.flash:\n self.register_buffer(\"bias\", torch.tril(torch.ones(config.block_size, config.block_size))\n .view(1, 1, config.block_size, config.block_size))\n \n self.use_pos_emb = False\n \n # RoPE frequencies (using base 100000 for better representation)\n inv_freq = 1.0 / (100000 ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))\n self.register_buffer(\"inv_freq\", inv_freq)\n \n # Token Shift mixing parameters for Q, K, V\n self.time_mix = nn.Parameter(torch.empty(1, 1, 3 * config.n_embd))\n nn.init.uniform_(self.time_mix, 0.2, 0.8)\n\n # Learnable attention temperature per head\n self.attn_scale = nn.Parameter(torch.ones(1, self.n_head, 1, 1))\n\n def _apply_rope(self, x, seq_len):\n t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)\n freqs = torch.outer(t, self.inv_freq)\n cos = freqs.cos().unsqueeze(0).unsqueeze(0)\n sin = freqs.sin().unsqueeze(0).unsqueeze(0)\n d = x.shape[-1] // 2\n x1, x2 = x[..., :d], x[..., d:]\n y1 = x1 * cos - x2 * sin\n y2 = x1 * sin + x2 * cos\n return torch.cat([y1, y2], dim=-1).type_as(x)\n\n def forward(self, x):\n B, T, C = x.size()\n \n qkv = self.c_attn(x)\n \n # Token Shift (Time Mixing)\n qkv_shifted = torch.cat([\n torch.zeros(B, 1, 3 * C, device=qkv.device, dtype=qkv.dtype), \n qkv[:, :-1, :]\n ], dim=1)\n qkv = qkv * self.time_mix + qkv_shifted * (1.0 - self.time_mix)\n \n q, k, v = qkv.split(self.n_embd, dim=2)\n \n k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)\n q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)\n v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)\n \n # QK-Norm without learned weights to preserve RoPE geometry\n q = F.rms_norm(q, (self.head_dim,))\n k = F.rms_norm(k, (self.head_dim,))\n \n # RoPE\n q = self._apply_rope(q, T)\n k = self._apply_rope(k, T)\n \n # Apply learnable temperature per head to Query\n q = q * self.attn_scale\n\n if self.flash:\n y = torch.nn.functional.scaled_dot_product_attention(\n q, k, v, attn_mask=None,\n dropout_p=self.dropout if self.training else 0, is_causal=True)\n else:\n att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))\n att = F.softmax(att, dim=-1)\n att = self.attn_dropout(att)\n y = att @ v\n \n y = y.transpose(1, 2).contiguous().view(B, T, C)\n y = self.resid_dropout(self.c_proj(y))\n return y"
8}Tool ResultStep 2
OK: Replaced lines 33..70 with 84 line(s) in nanoGPT/custom_pretrain.py. Editable range: 33–116, 291–293.
[Current file: nanoGPT/custom_pretrain.py | editable: 33–116 | total: 483 lines]
30| return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
31|
32| # ── Self-Attention ─────────────────────────────────────────────────────────
33 class CausalSelfAttention(nn.Module):
34 def __init__(self, config):
35 super().__init__()
36 assert config.n_embd % config.n_head == 0
37 self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
38 self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
39 self.attn_dropout = nn.Dropout(config.dropout)
40 self.resid_dropout = nn.Dropout(config.dropout)
41 self.n_head = config.n_head
42 self.n_embd = config.n_embd
43 self.head_dim = config.n_embd // co
AssistantStep 3test
Tool ResultStep 7
[Test #3] (0 tests remaining; call submit(n=N) to choose which test result to submit as final)
[NOTE] This was your last test. You MUST now call submit(n=X) to choose which test result to submit as your final answer.
### gpt-345m (scripts/gpt_345m.sh)
Training metrics (last steps):
TRAIN_METRICS: step=9000, train_loss=2.5908, val_loss=2.3145
TRAIN_METRICS: step=10000, train_loss=2.5613, val_loss=2.2849
TRAIN_METRICS: step=11000, train_loss=2.5466, val_loss=2.2675
TRAIN_METRICS: step=12000, train_loss=2.5311, val_loss=2.2628
TRAIN_METRICS: step=13000, train_loss=2.5216, val_loss=2.2509
Final metrics (gpt-345m): val_loss_gpt-345m=2.2564, wikitext2_ppl_gpt-345m=42.3800, lambada_ppl_gpt-345m=66.1500
### lm-eval-345m (scripts/lm_eval_345m.sh)
Final metrics (lm-eval-345m): arc_easy_lm-eval-345m=57.2800, hellaswag_lm-eval-345m=34.8900, piqa_lm-eval-345m=64.6900, winogrande_lm-eval-345m=52.1700
[Leaderboard] Results saved: {'val_loss_gpt-345m': 2.2564, 'wikitext2_ppl_gpt-345m': 42.38, 'lam
AssistantStep 8submit
Tool Input
1{
2 "n": 1
3}Tool ResultStep 8
[submit] Submitting result from test #1 as final.
### gpt-345m (scripts/gpt_345m.sh)
Training metrics (last steps):
TRAIN_METRICS: step=9000, train_loss=2.5886, val_loss=2.3115
TRAIN_METRICS: step=10000, train_loss=2.5594, val_loss=2.2825
TRAIN_METRICS: step=11000, train_loss=2.5446, val_loss=2.2650
TRAIN_METRICS: step=12000, train_loss=2.5292, val_loss=2.2606
TRAIN_METRICS: step=13000, train_loss=2.5199, val_loss=2.2489
Final metrics (gpt-345m): val_loss_gpt-345m=2.2552, wikitext2_ppl_gpt-345m=41.4500, lambada_ppl_gpt-345m=64.8000
### lm-eval-345m (scripts/lm_eval_345m.sh)
Final metrics (lm-eval-345m): arc_easy_lm-eval-345m=57.7000, hellaswag_lm-eval-345m=34.5500, piqa_lm-eval-345m=64.0400, winogrande_lm-eval-345m=50.4300
[submit] Finalized test #1 as final.
[Leaderboard] Results saved: {'val_loss_gpt-345m': 2.2552, 'wikitext2_ppl_gpt-345m': 41.45, 'lambada_ppl_gpt-345m': 64.8, 'elapsed_gpt-345m': 25558, 'arc_easy_lm-eval-345m': 57.7, 'hellaswag_lm-eval-345m': 34.55, 'piqa_lm-eval