MPNN and FFN dropout cannot be set independently; breaks parity with CheMeleon-initialized models

[smcolby]

Summary

ChemPropModel exposes a single dropout parameter that is applied to both the message-passing encoder and the FFN head. There is no way to set encoder (MPNN) dropout independently from FFN dropout. This makes it impossible to construct a from-scratch model that exactly matches the architecture of a CheMeleon-initialized one, because CheMeleon's encoder runs at dropout=0.0 while its FFN runs at the config's dropout.

Details

In the from-scratch branch, the single self.dropout feeds both the encoder and the FFN (openadmet/models/architecture/chemprop.py):

# encoder
mp = message_cls(
    d_h=self.message_hidden_dim, depth=self.depth, dropout=self.dropout   # L695-697
)
...
# head
ffn = nn.RegressionFFN(
    ...
    dropout=self.dropout,                                                 # L706
    ...
)

When initializing from CheMeleon, the encoder is instead built from the foundation's own hyperparameters and its state dict is loaded:

mp = nn.BondMessagePassing(**foundation_mp["hyper_parameters"])           # L676
mp.load_state_dict(foundation_mp["state_dict"])

The published CheMeleon checkpoint carries dropout=0.0 in hyper_parameters (verified from chemeleon_mp.pt: {'d_h': 2048, 'depth': 6, 'dropout': 0.0, 'bias': False}), so a CheMeleon model's effective dropout is encoder 0.0 / FFN = config dropout, whereas a from-scratch model is encoder = config / FFN = config.

Impact

For a controlled comparison isolating the effect of initialization (random vs CheMeleon weights) with architecture held fixed at CheMeleon's dims (message_hidden_dim=2048, depth=6), the only achievable settings are:

• dropout=0.1 → matches CheMeleon's FFN but not its encoder (encoder 0.1 vs 0.0)
• dropout=0.0 → matches CheMeleon's encoder but not its FFN (FFN 0.0 vs 0.1)

Neither reproduces CheMeleon's encoder-0.0 / FFN-0.1 split, so the from-scratch arm always differs from the CheMeleon arm by more than initialization alone.

Proposed fix

Keep dropout as the global, backward-compatible setting that drives both components when used alone. Introduce two optional overrides, mpnn_dropout: float | None and ffn_dropout: float | None, each defaulting to dropout when unset; when present, they take precedence for their respective component. This preserves all current behavior and lets callers express mpnn_dropout=0.0, ffn_dropout=0.1 for exact parity with CheMeleon. Include both new fields in the hparams-logged field list (L855-868) alongside dropout.