autoresearch

Score Breakdown

Your Submission

{
  "train_code": "import json\nimport math\nimport os\nimport time\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport prepare\nfrom prepare import BOS_TOKEN, MAX_SEQ_LEN, TIME_BUDGET, VOCAB_SIZE, evaluate_bpb\n\n\nchannels = 96\ndilations = [1, 2, 4, 8, 16, 32, 64, 128]\n\nbatch_size = 8\neval_batch_size = 64\nlearning_rate = 0.001\nmin_lr_ratio = 0.1\nwarmup_ratio = 0.06\nweight_decay = 0.01\ngrad_clip = 1.0\ntrain_deadline = TIME_BUDGET\ndevice = \"cpu\"\n\nseed = int(os.environ.get(\"TORCH_SEED\", \"42\"))\ntorch.manual_seed(seed)\n\n\ndef load_train_data() -> torch.Tensor:\n    shards = [prepare._load_shard(path) for path in prepare._get_data_files(\"train\")]\n    return torch.cat(shards)\n\n\ndef make_fast_batch(data: torch.Tensor, batch: int, seq_len: int, device_name: str):\n    starts = torch.randint(0, len(data) - seq_len - 1, (batch, 1))\n    offsets = torch.arange(seq_len, dtype=torch.long).unsqueeze(0)\n    targets = data[starts + offsets]\n    x = torch.empty((batch, seq_len), dtype=torch.long)\n    x[:, 0] = BOS_TOKEN\n    x[:, 1:] = targets[:, :-1]\n    return x.to(device_name), targets.to(device_name)\n\n\nclass CausalGLUBlock(nn.Module):\n    def __init__(self, width: int, dilation: int):\n        super().__init__()\n        self.norm = nn.LayerNorm(width)\n        self.conv = nn.Conv1d(width, 2 * width, kernel_size=3, dilation=dilation)\n        self.proj = nn.Conv1d(width, width, kernel_size=1)\n        self.pad = 2 * dilation\n\n    def forward(self, x):\n        residual = x\n        y = self.norm(x.transpose(1, 2)).transpose(1, 2)\n        y = self.conv(F.pad(y, (self.pad, 0)))\n        y = F.glu(y, dim=1)\n        y = self.proj(y)\n        return residual + y\n\n\nclass ByteCNN(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.tok_emb = nn.Embedding(VOCAB_SIZE, channels)\n        self.blocks = nn.ModuleList([CausalGLUBlock(channels, d) for d in dilations])\n        self.ln_f = nn.LayerNorm(channels)\n        self.head = nn.Linear(channels, VOCAB_SIZE, bias=False)\n        self.head.weight = self.tok_emb.weight\n        self.apply(self._init_weights)\n\n    @staticmethod\n    def _init_weights(module):\n        if isinstance(module, (nn.Linear, nn.Conv1d)):\n            nn.init.normal_(module.weight, mean=0.0, std=0.02)\n            if module.bias is not None:\n                nn.init.zeros_(module.bias)\n        elif isinstance(module, nn.Embedding):\n            nn.init.normal_(module.weight, mean=0.0, std=0.02)\n\n    def forward(self, x, targets=None):\n        h = self.tok_emb(x).transpose(1, 2)\n        for block in self.blocks:\n            h = block(h)\n        h = self.ln_f(h.transpose(1, 2))\n        logits = self.head(h)\n        if targets is None:\n            return logits\n        return F.cross_entropy(logits.reshape(-1, VOCAB_SIZE), targets.reshape(-1))\n\n\ndef build_optimizer(model: nn.Module) -> torch.optim.Optimizer:\n    decay_params = []\n    no_decay_params = []\n    for param in model.parameters():\n        if not param.requires_grad:\n            continue\n        if param.dim() >= 2:\n            decay_params.append(param)\n        else:\n            no_decay_params.append(param)\n    return torch.optim.AdamW(\n        [\n            {\"params\": decay_params, \"weight_decay\": weight_decay},\n            {\"params\": no_decay_params, \"weight_decay\": 0.0},\n        ],\n        lr=learning_rate,\n        betas=(0.9, 0.95),\n    )\n\n\ndef lr_for(elapsed: float) -> float:\n    progress = min(1.0, elapsed / train_deadline)\n    if progress < warmup_ratio:\n        return learning_rate * (progress / max(warmup_ratio, 1e-8))\n    cosine_progress = (progress - warmup_ratio) / max(1.0 - warmup_ratio, 1e-8)\n    cosine = 0.5 * (1.0 + math.cos(math.pi * cosine_progress))\n    return learning_rate * (min_lr_ratio + (1.0 - min_lr_ratio) * cosine)\n\n\ndef set_lr(optimizer: torch.optim.Optimizer, lr: float) -> None:\n    for group in optimizer.param_groups:\n        group[\"lr\"] = lr\n\n\ndef train():\n    train_data = load_train_data()\n    model = ByteCNN().to(device)\n    optimizer = build_optimizer(model)\n\n    start = time.time()\n    avg_step_time = 0.08\n    step = 0\n    last_loss = None\n\n    while True:\n        elapsed = time.time() - start\n        if elapsed + avg_step_time >= train_deadline:\n            break\n\n        set_lr(optimizer, lr_for(elapsed))\n        step_start = time.time()\n\n        x, y = make_fast_batch(train_data, batch_size, MAX_SEQ_LEN, device)\n        loss = model(x, y)\n        optimizer.zero_grad(set_to_none=True)\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n        optimizer.step()\n\n        step_time = time.time() - step_start\n        avg_step_time = 0.9 * avg_step_time + 0.1 * step_time\n        step += 1\n        last_loss = loss.item()\n\n    training_time = time.time() - start\n    eval_start = time.time()\n    val_bpb = evaluate_bpb(model, eval_batch_size, device)\n    eval_time = time.time() - eval_start\n\n    print(\n        json.dumps(\n            {\n                \"val_bpb\": round(val_bpb, 6),\n                \"train_loss\": round(last_loss, 6) if last_loss is not None else None,\n                \"total_steps\": step,\n                \"training_time_secs\": round(training_time, 2),\n                \"eval_time_secs\": round(eval_time, 2),\n                \"num_params_M\": round(sum(p.numel() for p in model.parameters()) / 1e6, 4),\n                \"batch_size\": batch_size,\n                \"eval_batch_size\": eval_batch_size,\n                \"learning_rate\": learning_rate,\n            }\n        )\n    )\n\n\nif __name__ == \"__main__\":\n    train()\n",
  "methodology": "## Experiment Log\n\n- run-0 `fastloader_transformer`: val_bpb=3.614273, steps=85, train_time=181.55s\n- run-1 `fastloader_short_context`: val_bpb=4.885838, steps=214, train_time=186.66s\n- run-2 `dilated_cnn`: val_bpb=3.142126, steps=83, train_time=186.38s\n- run-3 `smaller_cnn`: val_bpb=2.950446, steps=157, train_time=184.79s\n- run-4 `tuned_smaller_cnn`: val_bpb=2.915301, steps=206, train_time=189.79s\n- run-5 `smallest_batch_cnn`: val_bpb=2.885916, steps=246, train_time=189.88s\n- run-6 `final_cnn_tuning`: val_bpb=2.884851, steps=308, train_time=191.29s\n\n## Key Findings\n\n- Switching from the baseline transformer family to a causal dilated CNN was the main breakthrough. The CNN matched the fixed 256-byte evaluation context without quadratic attention cost on CPU.\n- A vectorized random byte batch loader was necessary, but loader speed alone was not enough; the architecture change was what moved val_bpb from the 3.5-4.9 range down into the 2.88-3.14 range.\n- The best model used 8 causal GLU residual blocks with dilations [1, 2, 4, 8, 16, 32, 64, 128], 96 channels, tied token/output weights, AdamW, cosine decay, and gradient clipping.\n- Within that CNN family, smaller training batches consistently helped. Moving from batch 24 to 16 to 8 increased update count from 157 to 206 to 308 and improved val_bpb from 2.950446 to 2.915301 to 2.884851.\n- This challenge remained compute-limited even after optimization: later runs still benefited from more updates, suggesting that additional improvement would likely come from either even more efficient kernels or a slightly better low-cost convolutional block.\n\n## Final Configuration\n\n- Model: causal dilated CNN, 96 channels, 8 residual GLU blocks, full 256-token receptive field\n- Optimizer: AdamW, lr=1.0e-3 peak, betas=(0.9, 0.95), weight_decay=0.01, grad clip=1.0\n- Batching: train batch size 8, eval batch size 64\n- Schedule: short warmup + cosine decay across the full training window\n\n## Best Result\n\n- Best run: `run-6`\n- val_bpb: 2.884851\n- total_steps: 308\n- training_time_secs: 191.29"
}