From 6dbe1f27b08119640837c4e903599628810baa5e Mon Sep 17 00:00:00 2001 From: alex-crlhmmr Date: Fri, 10 Jul 2026 04:13:02 -0700 Subject: [PATCH] Expose --attention_dropout in pretraining for consistency with downstream pretrain.py did not expose --attention_dropout, and models/pretraining.py built the graph encoder without passing it, so attention dropout stayed pinned at the DualGraphEncoder default (0.1) and --dropout did not affect it. classification.py and segmentation.py already expose the flag and thread it into DualGraphEncoder; this mirrors that pattern for pretraining. The default (0.1) preserves existing behavior; the flag now makes attention dropout controllable and visible in the pretraining config, which matters for reproducibility (the effective attention regularization was previously not derivable from the pretraining flags). --- models/pretraining.py | 1 + pretrain.py | 1 + 2 files changed, 2 insertions(+) diff --git a/models/pretraining.py b/models/pretraining.py index fa25741..e824e17 100644 --- a/models/pretraining.py +++ b/models/pretraining.py @@ -52,6 +52,7 @@ def __init__(self, args): num_layers=args.graph_num_layers, num_heads=args.graph_num_heads, dropout=args.dropout, + attention_dropout=getattr(args, "attention_dropout", args.dropout), dim_feedforward=args.dim_feedforward, add_positional_encoding=args.add_positional_encoding, use_edge_bias=args.use_edge_bias, diff --git a/pretrain.py b/pretrain.py index 372459d..e5d4767 100644 --- a/pretrain.py +++ b/pretrain.py @@ -26,6 +26,7 @@ def build_parser(): parser.add_argument("--dim_feedforward", type=int, default=512) parser.add_argument("--dropout", type=float, default=0.25) parser.add_argument("--mlp_dropout", type=float, default=0.1) + parser.add_argument("--attention_dropout", type=float, default=0.1) parser.add_argument("--act", type=str, default="gelu") parser.add_argument("--curve_emb_dim", type=int, default=64) parser.add_argument("--surface_emb_dim", type=int, default=64)