Add configurable dropout to ResBlocks

Wire a dropout hyperparameter (default 0.1) through the config, model,
training pipeline, and CLI. Persisted in saved model_config so checkpoints
reconstruct the architecture correctly.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

giant/config.py

@@ -14,7 +14,7 @@ DEFAULT_CONFIG: dict = {
         "val_fraction": 0.1, "num_workers": 4, "seed": 0, "validate_every": 10,
     },
     "model": {
-        "hidden_dim": 256, "n_blocks": 6, "emb_dim": 16,
+        "hidden_dim": 256, "n_blocks": 6, "emb_dim": 16, "dropout": 0.1,
     },
 }
 

giant/model/network.py

@@ -49,18 +49,20 @@ class ConditionEncoder(nn.Module):
 
 
 class ResBlock(nn.Module):
-    def __init__(self, dim: int, cond_dim: int) -> None:
+    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.1) -> None:
         super().__init__()
         self.norm = nn.LayerNorm(dim)
         self.linear1 = nn.Linear(dim, dim)
         self.cond_proj = nn.Linear(cond_dim, dim, bias=False)
         self.act = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
         self.linear2 = nn.Linear(dim, dim)
 
     def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
         h = self.norm(x)
         h = self.linear1(h) + self.cond_proj(cond)
         h = self.act(h)
+        h = self.dropout(h)
         h = self.linear2(h)
         return x + h
 
@@ -76,6 +78,7 @@ class DenoisingMLP(nn.Module):
         time_dim: int = 64,
         cond_out_dim: int = 128,
         x_dim: int = X_DIM,
+        dropout: float = 0.1,
     ) -> None:
         super().__init__()
         self.time_emb = SinusoidalEmbedding(time_dim)
@@ -88,7 +91,7 @@ class DenoisingMLP(nn.Module):
         merged_cond_dim = time_dim + cond_out_dim
         self.input_proj = nn.Linear(x_dim, hidden_dim)
         self.blocks = nn.ModuleList([
-            ResBlock(hidden_dim, merged_cond_dim) for _ in range(n_blocks)
+            ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)
         ])
         self.out_proj = nn.Linear(hidden_dim, x_dim)
 

giant/pipeline.py

@@ -96,6 +96,7 @@ def run_train_job(
     model = DenoisingMLP(
         pdg_vocab=len(pdg_map), mat_vocab=len(mat_map),
         hidden_dim=m["hidden_dim"], n_blocks=m["n_blocks"], emb_dim=m["emb_dim"],
+        dropout=m["dropout"],
     )
     echo(f"model: {sum(p.numel() for p in model.parameters()):,} parameters")
 
@@ -114,6 +115,7 @@ def run_train_job(
     model_config = {
         "pdg_vocab": len(pdg_map), "mat_vocab": len(mat_map),
         "hidden_dim": m["hidden_dim"], "n_blocks": m["n_blocks"], "emb_dim": m["emb_dim"],
+        "dropout": m["dropout"],
     }
 
     run_training(

scripts/train.py

@@ -18,6 +18,7 @@ def main() -> None:
     parser.add_argument("--hidden-dim", type=int)
     parser.add_argument("--n-blocks", type=int)
     parser.add_argument("--emb-dim", type=int)
+    parser.add_argument("--dropout", type=float, help="Dropout probability in ResBlocks (default: 0.1)")
     parser.add_argument("--val-fraction", type=float)
     parser.add_argument("--seed", type=int, help="Random seed for reproducibility")
     parser.add_argument("--validate-every", type=int,
@@ -37,6 +38,7 @@ def main() -> None:
     }.items() if v is not None}
     cli_model = {k: v for k, v in {
         "hidden_dim": args.hidden_dim, "n_blocks": args.n_blocks, "emb_dim": args.emb_dim,
+        "dropout": args.dropout,
     }.items() if v is not None}
     config_path = Path(args.config) if args.config else None
     cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config_path, cli_train, cli_model)