scripts/data_scaling_full.py

#!/usr/bin/env python3
"""Full data scaling sweep for Claim 6.
Trains toy DiT with N=20,50,100,200,400,800 samples.
Paper: 2512.24165
"""
import modal, json, os, time, sys
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

app = modal.App("diffthinker-data-scaling")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("torch>=2.1.0", "torchvision", "numpy", "Pillow")
)

def make_maze_data(grid_size=8, num=100):
    import torch
    data = []
    for _ in range(num):
        g = np.zeros((grid_size, grid_size), dtype=int)
        start = (0, np.random.randint(0, grid_size))
        goal = (grid_size-1, np.random.randint(0, grid_size))
        for _ in range(int(grid_size*grid_size*0.15)):
            wx, wy = np.random.randint(1, grid_size-1, 2)
            g[wx, wy] = 1
        from collections import deque
        q = deque([start])
        parent = {start: None}
        solved = False
        while q:
            r, c = q.popleft()
            if (r, c) == goal:
                solved = True; break
            for dr, dc in [(0,1),(0,-1),(1,0),(-1,0)]:
                nr, nc = r+dr, c+dc
                if 0 <= nr < grid_size and 0 <= nc < grid_size and g[nr,nc] != 1 and (nr,nc) not in parent:
                    parent[(nr,nc)] = (r,c)
        # Build solution path
        path_grid = np.zeros((grid_size, grid_size), dtype=np.float32)
        if solved:
            cur = goal
            while cur != start:
                path_grid[cur] = 1.0
                cur = parent[cur]
            path_grid[start] = 1.0
        # Input: walls+start+goal
        inp = np.stack([g==1, np.zeros_like(g), np.zeros_like(g)], axis=0).astype(np.float32)
        inp[1, start[0], start[1]] = 1.0
        inp[2, goal[0], goal[1]] = 1.0
        data.append({"input": inp, "target": path_grid, "solved": solved})
    return data

class TinyDiT(nn.Module):
    def __init__(self, grid_size=8, hidden=64):
        super().__init__()
        self.grid_size = grid_size
        self.hidden = hidden
        self.time_embed = nn.Sequential(
            nn.Linear(1, hidden), nn.SiLU(), nn.Linear(hidden, hidden))
        self.input_proj = nn.Conv2d(3, hidden, 3, padding=1)
        self.block = nn.Sequential(
            nn.Conv2d(hidden+hidden, hidden, 3, padding=1), nn.SiLU(),
            nn.Conv2d(hidden, hidden, 3, padding=1), nn.SiLU())
        self.out = nn.Conv2d(hidden, 1, 1)

    def forward(self, x, t):
        B = x.shape[0]
        te = self.time_embed(t.view(B,1))
        h = self.input_proj(x)
        te = te.view(B, self.hidden, 1, 1).expand(-1, -1, self.grid_size, self.grid_size)
        h = torch.cat([h, te], dim=1)
        h = self.block(h)
        return self.out(h)

@app.cls(image=image, gpu="T4", scaledown_window=60, timeout=1800)
class DataScaling:
    def train_and_eval(self, num_train):
        device = "cuda"
        data = make_maze_data(num=num_train)
        model = TinyDiT().to(device)
        opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
        N = len(data)
        losses = []
        for epoch in range(30):
            epoch_loss = 0.0
            idxs = np.random.permutation(N)
            for i in idxs:
                inp = torch.tensor(data[i]["input"]).unsqueeze(0).to(device)
                target = torch.tensor(data[i]["target"]).unsqueeze(0).unsqueeze(0).to(device)
                t = torch.rand(1,1).to(device)
                noise = torch.randn_like(target)
                v_pred = model(inp, t)
                v_target = noise - target
                loss = F.mse_loss(v_pred, v_target)
                opt.zero_grad()
                loss.backward()
                opt.step()
                epoch_loss += loss.item()
            avg_loss = epoch_loss / N
            losses.append(avg_loss)
            if (epoch+1) % 5 == 0:
                print(f"  Ep {epoch+1}/30 loss={avg_loss:.4f}")

        eval_data = make_maze_data(num=50)
        correct = 0
        for d in eval_data:
            inp = torch.tensor(d["input"]).unsqueeze(0).to(device)
            x = torch.randn(1,1,8,8).to(device)
            steps = 20
            for s in range(steps):
                t = torch.tensor([[(s+1)/steps]]).to(device)
                v = model(inp, t)
                x = x + v * (1.0/steps)
            pred_path = (x.squeeze().detach().cpu().numpy() > 0.5).astype(int)
            gt = d["target"]
            acc = (pred_path == gt).mean()
            correct += 1 if acc > 0.9 else 0
        acc = correct / len(eval_data) * 100
        print(f"  Acc: {acc:.1f}% ({correct}/{len(eval_data)})")
        return {"num_train": num_train, "final_loss": round(losses[-1], 4), "accuracy": round(acc, 1), "epochs": 30}

    @modal.method()
    def sweep(self, data_sizes=[20, 50, 100, 200, 400, 800]):
        results = []
        for n in data_sizes:
            print(f"\n=== Training with N={n} samples ===")
            t0 = time.time()
            r = self.train_and_eval(n)
            r["time_sec"] = round(time.time()-t0, 1)
            results.append(r)
            print(f"   Done in {r['time_sec']:.1f}s")
        print(json.dumps(results, indent=2))
        return results

@app.local_entrypoint()
def main():
    ds = DataScaling()
    res = ds.sweep.remote([20, 50, 100, 200, 400, 800])
    with open("/tmp/data_scaling_results.json", "w") as f:
        json.dump(res, f, indent=2)
    print("Results saved to /tmp/data_scaling_results.json")