From da848016257cdb129703b2ba0a855dcf2be4b627 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 18 Jun 2026 15:59:20 +0200 Subject: [PATCH] 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 --- giant/config.py | 2 +- giant/model/network.py | 7 +++++-- giant/pipeline.py | 2 ++ scripts/train.py | 2 ++ 4 files changed, 10 insertions(+), 3 deletions(-) diff --git a/giant/config.py b/giant/config.py index 2b413a3..f861ad8 100644 --- a/giant/config.py +++ b/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, }, } diff --git a/giant/model/network.py b/giant/model/network.py index 9ff099c..e619fdb 100644 --- a/giant/model/network.py +++ b/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) diff --git a/giant/pipeline.py b/giant/pipeline.py index 1e5c8d2..5909c05 100644 --- a/giant/pipeline.py +++ b/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( diff --git a/scripts/train.py b/scripts/train.py index 9a5bae7..84b3540 100644 --- a/scripts/train.py +++ b/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)