fix: proper controls — loss masking, uniform baseline, multi-seed

src/fog/config.py

@@ -34,6 +34,17 @@ MOTIF_SMALL = FOGConfig(
     d_gate=32,
 )
 
+# Param-matched uniform baseline for controlled comparison
+# d_model=94, d_ff=376 → ~432K params to match MOTIF_TINY
+UNIFORM_TINY = FOGConfig(
+    vocab_size=32,
+    d_model=94,
+    n_layers=4,
+    n_heads=2,
+    max_seq_len=32,
+    d_ff=376,
+)
+
 # Tiny configs for fast iteration
 BASELINE_TINY = FOGConfig(
     vocab_size=32,

src/fog/data.py

@@ -16,17 +16,18 @@ class CopyTask(Dataset):
         self.sep_token = vocab_size - 1
         rng = random.Random(seed)
         self.samples = []
-        content_vocab = vocab_size - 1  # exclude SEP
+        self.sep_positions = []
+        content_vocab = vocab_size - 1
         half = seq_len // 2 - 1
         for _ in range(n_samples):
             content = [rng.randint(0, content_vocab - 1) for _ in range(half)]
-            # input: content + SEP + content (teacher forcing)
             ids = content + [self.sep_token] + content
-            # pad/truncate to seq_len
+            sep_pos = len(content)
             ids = ids[:seq_len]
             while len(ids) < seq_len:
                 ids.append(0)
             self.samples.append(ids)
+            self.sep_positions.append(sep_pos)
 
     def __len__(self) -> int:
         return len(self.samples)
@@ -35,7 +36,12 @@ class CopyTask(Dataset):
         ids = self.samples[idx]
         x = torch.tensor(ids[:-1], dtype=torch.long)
         y = torch.tensor(ids[1:], dtype=torch.long)
-        return {"input_ids": x, "targets": y}
+        # loss_mask: 1 after SEP, 0 before (shifted by -1 for targets)
+        mask = torch.zeros_like(y)
+        sep = self.sep_positions[idx]
+        if sep < len(mask):
+            mask[sep:] = 1
+        return {"input_ids": x, "targets": y, "loss_mask": mask}
 
 
 class ReverseTask(Dataset):
@@ -47,15 +53,18 @@ class ReverseTask(Dataset):
         self.sep_token = vocab_size - 1
         rng = random.Random(seed)
         self.samples = []
+        self.sep_positions = []
         content_vocab = vocab_size - 1
         half = seq_len // 2 - 1
         for _ in range(n_samples):
             content = [rng.randint(0, content_vocab - 1) for _ in range(half)]
             ids = content + [self.sep_token] + list(reversed(content))
+            sep_pos = len(content)
             ids = ids[:seq_len]
             while len(ids) < seq_len:
                 ids.append(0)
             self.samples.append(ids)
+            self.sep_positions.append(sep_pos)
 
     def __len__(self) -> int:
         return len(self.samples)
@@ -64,7 +73,11 @@ class ReverseTask(Dataset):
         ids = self.samples[idx]
         x = torch.tensor(ids[:-1], dtype=torch.long)
         y = torch.tensor(ids[1:], dtype=torch.long)
-        return {"input_ids": x, "targets": y}
+        mask = torch.zeros_like(y)
+        sep = self.sep_positions[idx]
+        if sep < len(mask):
+            mask[sep:] = 1
+        return {"input_ids": x, "targets": y, "loss_mask": mask}
 
 
 class SelectiveRetrieval(Dataset):
@@ -79,6 +92,7 @@ class SelectiveRetrieval(Dataset):
         self.sep_token = vocab_size - 1
         rng = random.Random(seed)
         self.samples = []
+        self.sep_positions = []
         content_vocab = vocab_size - 2  # exclude SEP and padding
         for _ in range(n_samples):
             keys = rng.sample(range(content_vocab), min(n_pairs, content_vocab))
@@ -88,6 +102,7 @@ class SelectiveRetrieval(Dataset):
             ids = []
             for k, v in zip(keys, values):
                 ids.extend([k, v])
+            sep_pos = len(ids)
             ids.append(self.sep_token)
             ids.append(keys[query_idx])
             ids.append(values[query_idx])
@@ -96,6 +111,7 @@ class SelectiveRetrieval(Dataset):
             while len(ids) < seq_len:
                 ids.append(0)
             self.samples.append(ids)
+            self.sep_positions.append(sep_pos)
 
     def __len__(self) -> int:
         return len(self.samples)
@@ -104,4 +120,8 @@ class SelectiveRetrieval(Dataset):
         ids = self.samples[idx]
         x = torch.tensor(ids[:-1], dtype=torch.long)
         y = torch.tensor(ids[1:], dtype=torch.long)
-        return {"input_ids": x, "targets": y}
+        mask = torch.zeros_like(y)
+        sep = self.sep_positions[idx]
+        if sep < len(mask):
+            mask[sep:] = 1
+        return {"input_ids": x, "targets": y, "loss_mask": mask}

src/fog/model_baseline.py

@@ -79,6 +79,7 @@ class BaselineTransformer(nn.Module):
         self,
         input_ids: torch.Tensor,
         targets: torch.Tensor | None = None,
+        loss_mask: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         b, t = input_ids.shape
         pos = torch.arange(t, device=input_ids.device).unsqueeze(0)
@@ -95,6 +96,16 @@ class BaselineTransformer(nn.Module):
 
         loss = None
         if targets is not None:
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+            if loss_mask is not None:
+                # only compute loss on target positions (after SEP)
+                flat_logits = logits.view(-1, logits.size(-1))
+                flat_targets = targets.view(-1)
+                flat_mask = loss_mask.view(-1).bool()
+                if flat_mask.any():
+                    loss = F.cross_entropy(flat_logits[flat_mask], flat_targets[flat_mask])
+                else:
+                    loss = torch.tensor(0.0, device=logits.device)
+            else:
+                loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
 
         return {"logits": logits, "loss": loss}

src/fog/model_motif.py

@@ -121,6 +121,7 @@ class MotifTransformer(nn.Module):
         self,
         input_ids: torch.Tensor,
         targets: torch.Tensor | None = None,
+        loss_mask: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         b, t = input_ids.shape
         pos = torch.arange(t, device=input_ids.device).unsqueeze(0)
@@ -136,6 +137,15 @@ class MotifTransformer(nn.Module):
 
         loss = None
         if targets is not None:
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+            if loss_mask is not None:
+                flat_logits = logits.view(-1, logits.size(-1))
+                flat_targets = targets.view(-1)
+                flat_mask = loss_mask.view(-1).bool()
+                if flat_mask.any():
+                    loss = F.cross_entropy(flat_logits[flat_mask], flat_targets[flat_mask])
+                else:
+                    loss = torch.tensor(0.0, device=logits.device)
+            else:
+                loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
 
         return {"logits": logits, "loss": loss}

src/fog/train.py

@@ -9,7 +9,7 @@ from pathlib import Path
 import torch
 from torch.utils.data import DataLoader
 
-from src.fog.config import FOGConfig, BASELINE_SMALL, MOTIF_SMALL, BASELINE_TINY, MOTIF_TINY
+from src.fog.config import FOGConfig, BASELINE_SMALL, MOTIF_SMALL, BASELINE_TINY, MOTIF_TINY, UNIFORM_TINY
 from src.fog.model_baseline import BaselineTransformer
 from src.fog.model_motif import MotifTransformer
 from src.fog.data import CopyTask, ReverseTask, SelectiveRetrieval
@@ -31,7 +31,10 @@ def train_epoch(
     for batch in loader:
         input_ids = batch["input_ids"].to(device)
         targets = batch["targets"].to(device)
-        out = model(input_ids, targets)
+        loss_mask = batch.get("loss_mask")
+        if loss_mask is not None:
+            loss_mask = loss_mask.to(device)
+        out = model(input_ids, targets, loss_mask=loss_mask)
         loss = out["loss"]
         optimizer.zero_grad()
         loss.backward()
@@ -58,21 +61,29 @@ def eval_accuracy(
     for batch in loader:
         input_ids = batch["input_ids"].to(device)
         targets = batch["targets"].to(device)
-        out = model(input_ids, targets)
+        loss_mask = batch.get("loss_mask")
+        if loss_mask is not None:
+            loss_mask = loss_mask.to(device)
+        out = model(input_ids, targets, loss_mask=loss_mask)
         total_loss += out["loss"].item()
         n_batches += 1
 
         preds = out["logits"].argmax(dim=-1)
-        # only count accuracy after SEP token
-        for i in range(input_ids.size(0)):
-            sep_positions = (input_ids[i] == sep_token).nonzero(as_tuple=True)[0]
-            if len(sep_positions) == 0:
-                continue
-            start = sep_positions[0].item() + 1
-            if start >= targets.size(1):
-                continue
-            correct += (preds[i, start:] == targets[i, start:]).sum().item()
-            total += targets.size(1) - start
+        # accuracy only on masked (target) positions
+        if loss_mask is not None:
+            m = loss_mask.bool()
+            correct += (preds[m] == targets[m]).sum().item()
+            total += m.sum().item()
+        else:
+            for i in range(input_ids.size(0)):
+                sep_positions = (input_ids[i] == sep_token).nonzero(as_tuple=True)[0]
+                if len(sep_positions) == 0:
+                    continue
+                start = sep_positions[0].item() + 1
+                if start >= targets.size(1):
+                    continue
+                correct += (preds[i, start:] == targets[i, start:]).sum().item()
+                total += targets.size(1) - start
 
     return {
         "loss": total_loss / max(n_batches, 1),
@@ -89,17 +100,20 @@ def run_experiment(
     batch_size: int,
     lr: float,
     device: torch.device,
+    seed: int = 42,
 ) -> dict:
-    # Data
+    torch.manual_seed(seed)
+
+    # Data — use fixed seeds for data, model seed varies
     n_train, n_eval = 5000, 500
     if task_name == "copy":
-        train_ds = CopyTask(cfg.vocab_size, cfg.max_seq_len, n_train, seed=42)
+        train_ds = CopyTask(cfg.vocab_size, cfg.max_seq_len, n_train, seed=0)
         eval_ds = CopyTask(cfg.vocab_size, cfg.max_seq_len, n_eval, seed=99)
     elif task_name == "reverse":
-        train_ds = ReverseTask(cfg.vocab_size, cfg.max_seq_len, n_train, seed=42)
+        train_ds = ReverseTask(cfg.vocab_size, cfg.max_seq_len, n_train, seed=0)
         eval_ds = ReverseTask(cfg.vocab_size, cfg.max_seq_len, n_eval, seed=99)
     elif task_name == "retrieval":
-        train_ds = SelectiveRetrieval(cfg.vocab_size, cfg.max_seq_len, n_train, seed=42)
+        train_ds = SelectiveRetrieval(cfg.vocab_size, cfg.max_seq_len, n_train, seed=0)
         eval_ds = SelectiveRetrieval(cfg.vocab_size, cfg.max_seq_len, n_eval, seed=99)
     else:
         raise ValueError(f"Unknown task: {task_name}")
@@ -141,6 +155,7 @@ def run_experiment(
     return {
         "model_type": model_type,
         "task": task_name,
+        "seed": seed,
         "n_params": n_params,
         "n_epochs": n_epochs,
         "elapsed_s": round(elapsed, 1),
@@ -159,37 +174,40 @@ def main() -> None:
     parser.add_argument("--lr", type=float, default=3e-4)
     parser.add_argument("--device", type=str, default="cpu")
     parser.add_argument("--size", type=str, default="tiny", choices=["tiny", "small"])
+    parser.add_argument("--seeds", type=int, nargs="+", default=[42])
     parser.add_argument("--output", type=str, default="archive/fog_ablation.json")
     args = parser.parse_args()
 
     device = torch.device(args.device)
 
     if args.size == "tiny":
-        configs = [("baseline", BASELINE_TINY), ("motif", MOTIF_TINY)]
+        configs = [("baseline", BASELINE_TINY), ("uniform_small", UNIFORM_TINY), ("motif", MOTIF_TINY)]
     else:
         configs = [("baseline", BASELINE_SMALL), ("motif", MOTIF_SMALL)]
 
     results = []
 
     for task in args.tasks:
-        print(f"\n{'='*60}")
-        print(f"  Task: {task} (size={args.size})")
-        print(f"{'='*60}")
-
-        for model_type, cfg in configs:
-            result = run_experiment(
-                task_name=task,
-                cfg=cfg,
-                model_type=model_type,
-                n_epochs=args.epochs,
-                batch_size=args.batch_size,
-                lr=args.lr,
-                device=device,
-            )
-            results.append(result)
-            print(f"  → {model_type}: params={result['n_params']:,}  "
-                  f"acc={result['final_accuracy']:.4f}  "
-                  f"time={result['elapsed_s']}s")
+        for seed in args.seeds:
+            print(f"\n{'='*60}")
+            print(f"  Task: {task} (size={args.size}, seed={seed})")
+            print(f"{'='*60}")
+
+            for model_type, cfg in configs:
+                result = run_experiment(
+                    task_name=task,
+                    cfg=cfg,
+                    model_type=model_type,
+                    n_epochs=args.epochs,
+                    batch_size=args.batch_size,
+                    lr=args.lr,
+                    device=device,
+                    seed=seed,
+                )
+                results.append(result)
+                print(f"  -> {model_type}: params={result['n_params']:,}  "
+                      f"acc={result['final_accuracy']:.4f}  "
+                      f"time={result['elapsed_s']}s")
 
     # Summary
     print(f"\n{'='*60}")