End-to-end Chemprop + Mother MLP/Flow Heads

[thomasATbayer]

End-to-end Chemprop + Mother MLP/Flow Heads

Problem

Current Chemprop workflows in Mother are strongest when embeddings are precomputed and then modeled downstream. That is useful, but it prevents representation learning from adapting to the final head objective during training.

We want a true end-to-end path where Chemprop and the Mother head are optimized jointly in one differentiable model.

Proposal

Add end-to-end estimators that combine a Chemprop encoder with Mother neural heads:

• ChempropMLPRegressor
• ChempropMLPClassifier
• ChempropFlowRegressor

Target training graph:

SMILES or molecular graph -> Chemprop encoder -> latent embedding -> Mother head -> task loss

Key requirement: gradients from the head must update Chemprop encoder weights.

Also required: users should be able to initialize the Chemprop encoder from
pretrained CheMeleon weights before end-to-end fine-tuning.

Scope

In scope:

• End-to-end regression and classification with MLP heads
• End-to-end probabilistic regression with Flow head
• Uncertainty outputs through the same Mother schema
• Quantiles for flow regression

Out of scope:

• Tree backbones
• Multi-backbone ensembling
• Export/distillation optimization work

Proposed API

Each estimator should expose:

• fit(X, y)
• predict(X)
• predict_uncertainty(X, ...)
• predict_quantiles(X, quantiles=...) for flow regressor only

Input X should accept Chemprop-compatible molecular inputs (at minimum SMILES).

Constructor parameters should be separated into:

• encoder_* for Chemprop architecture
• head_* for MLP/flow architecture
• training params such as lr, batch_size, max_epochs, callbacks

The API should include an explicit initialization path for pretrained CheMeleon
weights, for example via encoder_pretrained="chemeleon" and/or
encoder_checkpoint_path=....

Example shape:

ChempropFlowRegressor(
  encoder_hidden_dim=300,
  encoder_depth=3,
  encoder_dropout=0.1,
  encoder_pretrained="chemeleon",
  encoder_checkpoint_path=None,  # optional explicit checkpoint override
  head_mlp_hidden_dims="auto",
  head_mlp_dropout=0.1,
  flow_type="NSF",
  flow_transforms=3,
  max_epochs=100,
  lr=1e-3,
)

Implementation milestones

1. Build a Chemprop encoder wrapper module that outputs a dense embedding tensor and exposes embedding_dim.

• Include checkpoint-loading logic for pretrained CheMeleon weights.
• Support both default CheMeleon initialization and explicit checkpoint path.

2. Add joint modules:
   • Chemprop + MLP head
   • Chemprop + Flow head
3. Add estimator wrappers with sklearn-compatible parameter handling and cloning behavior.
4. Wire task losses:
   • MLP regressor: MSE
   • MLP classifier: CrossEntropy or BCEWithLogits
   • Flow regressor: NLL via log_prob
5. Integrate uncertainty paths:
   • MLP heads: MC-dropout uncertainty
   • Flow head: data/knowledge/total decomposition, plus quantiles
6. Add data utilities and robust error handling for malformed molecular inputs.
7. Add docs and examples for end-to-end MLP and Flow usage.

Acceptance criteria

• End-to-end estimators train on a small molecular dataset without manual representation precomputation.
• End-to-end estimators support initialization from pretrained CheMeleon weights.
• predict and predict_uncertainty work for all new estimators.
• predict_quantiles works for the flow estimator.
• Returned uncertainty matches Mother output schema.
• sklearn clone and CV compatibility is preserved.
• Unit tests cover forward pass, loss routing, uncertainty invariants, and clone round-trip.

Example usage target

# End-to-end Chemprop + MLP regression
reg = ChempropMLPRegressor(
  encoder_hidden_dim=300,
  encoder_depth=3,
  head_hidden_dims=[256, 128, 64],
  head_dropout=0.05,
)
reg.fit(smiles_train, y_train)
unc = reg.predict_uncertainty(smiles_test)

# End-to-end Chemprop + Flow probabilistic regression
flow = ChempropFlowRegressor(
  encoder_hidden_dim=300,
  encoder_depth=3,
  head_mlp_hidden_dims="auto",
  head_mlp_dropout=0.1,
  flow_type="NSF",
)
flow.fit(smiles_train, y_train)
unc = flow.predict_uncertainty(smiles_test)
q = flow.predict_quantiles(smiles_test, quantiles=[0.1, 0.5, 0.9])

Open questions

• Should v1 support both SMILES and prebuilt graph objects, or start with SMILES only?
• Do we want an optional freeze-then-unfreeze encoder schedule?
• Should mixed precision be enabled by default for larger runs?