GRASP-A2

GRASP-A2 adds explicit bond and stereochemistry inputs to the position-free GRASP-A1 encoder. This release contains the completed Step 2 encoder after 50,000 ChEMBL adaptation updates. It has 93,098,712 trainable encoder parameters, 12 Transformer layers, hidden size 768 and 12 attention heads: the same full-size architecture family as GRASP, rather than GRASP-Small.

A2 is a research ablation. It shows a small, uneven OpenADMET gain in the matched first training replicate; repeat-stable or universal superiority has not been established.

Install and encode molecules

Use Python 3.10–3.13. The release includes an installable, source-containing loader wheel with A2's chemistry-aware input pipeline. The older GRASP loader does not support these inputs. RDKit 2025.3.6 is pinned by the wheel to preserve the feature schema.

python -m pip install "https://huggingface.co/caithmac/GRASP-A2/resolve/main/grasp_molecular-0.1.1-py3-none-any.whl"

Restart a notebook runtime if an older grasp package was already imported. Then:

from grasp import GRASPEncoder

encoder = GRASPEncoder.from_pretrained("caithmac/GRASP-A2")
vectors = encoder.encode(
    ["CCO", "F[C@H](Cl)Br", "F[C@@H](Cl)Br"],
    batch_size=3,
)
print(vectors.shape)  # (3, 768)

The loader selects CUDA when available; pass device="cpu" for CPU inference. encode returns final-layer CLS embeddings, not predicted ADMET properties or affinities. The 23 endpoint-specific fine-tuned heads are not included. A standard Transformers AutoModel interface is not provided.

The model parses SMILES as supplied, supports up to 511 atoms plus CLS, and rejects invalid inputs. It does not silently strip salts, neutralize charges, choose tautomers, or invent unspecified stereochemistry.

Fine-tune on your task

Prepare separate CSVs with smiles and target columns:

from grasp.train import fit

fit(
    model_source="caithmac/GRASP-A2",
    train_csv="train.csv",
    valid_csv="valid.csv",
    output_dir="a2_property_model",
    task="regression",       # or "binary" with targets 0/1
    method="full",           # also "frozen" or "lora"
    epochs=60,
    patience=12,
    batch_size=64,
    learning_rate=1e-5,
    seed=0,
)

This uses learned mixing of layers 4/8/10/12 and atom-attention pooling. Frozen adaptation still trains the pooler and prediction head; LoRA adds query/value adapters. Choose settings using validation data, keep molecular groups separated between partitions, and evaluate the held-out test only after selection. Smaller GPU memory may require a smaller batch; that changes the downstream optimization recipe. A2-specific LoRA efficacy is not established by the compatibility tests in this release.

from grasp import GRASPPredictor

predictor = GRASPPredictor.from_pretrained("a2_property_model")
predictions = predictor.predict(["CCO", "CC(=O)O"])

What changed from GRASP and A1?

Component	Original GRASP architecture	A1	A2
Absolute atom positions	Learned	Removed	Removed
Shortest-path relative attention	Yes	Yes	Yes
Radius-0 atom-environment vocabulary	Yes	Yes	Yes
Explicit bond type/aromaticity inputs	No	No	Yes
Explicit atom/bond stereochemistry inputs	No	No	Yes

A2 adds atom-stereo embeddings and per-head bond attention biases. Atom categories include R/S, r/s, stereogenic-but-unspecified, nonstereogenic, unknown and unsupported cases. Bond categories distinguish E/Z, unspecified, nonstereogenic, unknown and unsupported cases. Bond inputs also distinguish single/double/triple/aromatic bonds and directional dative bonds. The precise version-1 schema and extraction implementation are included in the loader wheel.

These inputs let specified stereoisomers produce different model inputs. This is not proof that every stereochemical activity effect has been learned. A2 retains the original atom-token objective; factorized atom features and their corresponding objective belong to the separate A3 experiment.

Training provenance

• Structural Step 1: graph replaced-token detection with GDES, using a scheduled 1,538,242,560 molecule presentations from ZINC20. Presentations do not certify unique-molecule coverage. The original atom-environment vocabulary is retained.
• Step 2: 50,000 updates on an audited 511,897-molecule ChEMBL set with 642 assay tasks; effective batch 512, learning rate 5e-6 and seed 42. The learning-rate schedule retains its original 80,000-step horizon.
• Released checkpoint: the final Step 2 encoder used in A2's completed OpenADMET evaluations. This repository does not include the Step 1 checkpoint or optimizer/training states.
• Source checkpoint SHA256: 15da88b4c3e079b127142cc7bd0eb61819fd94d3aba3344ca919f659876c023f.

See training_provenance.json, config.json, validation_report.json and SHA256SUMS.

OpenADMET results

All models below use replicate 0 on three identical stored cluster-held-out splits for each of 23 endpoints. All use full fine-tuning with learned layers 4/8/10/12, batch 64, learning rate 1e-5, maximum 60 epochs and patience 12. Cohort hashes, ordered targets, seeds, initialized heads and scientific settings were checked; all 207 individual model scores were recomputed from saved predictions.

Model	Mean endpoint MAE ↓	Lowest mean among the three models
Matched original-architecture GRASP control	0.370232	8/23 endpoints
A1	0.368588	10/23 endpoints
A2	0.366232	5/23 endpoints

A2 reduces macro MAE by 1.08% versus control, with lower means on 14/23 endpoints. Versus A1 the reduction is 0.64%, but A2 improves only 8/23 endpoints: fewer larger gains outweigh more numerous smaller losses.

Exploratory endpoint-bootstrap 95% intervals for the mean paired MAE difference are [-0.012636, +0.004426] for A2 minus control and [-0.010505, +0.004860] for A2 minus A1. Both include zero. These intervals condition on the fitted models, omit pretraining-seed variability, and have no multiplicity correction. The endpoints share source collections. Each endpoint retains its own target scale, so the macro MAE is collection-specific.

The matched control above is a different checkpoint from the older released GRASP behind the paper table, whose historical mean MAE was 0.374183. It must not be silently substituted for that historical release. The completed control/A1 repeats average to 0.365261/0.369438; A2 has no corresponding additional repeats, so a balanced repeated comparison is pending. TDC results for A2 are also pending. The unrelated six-dataset Colab scaffold holdouts are not used here.

Positive values indicate lower error with A2. All 23 endpoints, split SDs and percentage changes are in the comparison CSV.

Validation and limitations

Release checks cover tensor-identical conversion to safetensors, strict loading, EIDF input/output parity, specified enantiomer and E/Z controls, randomized-SMILES controls, and tiny CPU regression/binary adaptation and save/reload tests. Legacy loader tests and A2-specific feature/gradient tests also passed. Tiny fits verify execution, not predictive efficacy; only the completed OpenADMET experiment above supplies benchmark results.

The encoder uses 2D graph information and supplied stereochemistry, without explicit solvent, pH, conformer ensembles, protein context or quantum-mechanical supervision. Unsupported stereo types use explicit fallbacks. No chemical mechanism for individual benchmark gains or losses has been demonstrated. One pretrained checkpoint per architecture limits inference, and pretraining/downstream overlap has not been certified absent for this release. Test-time atom ordering may produce small floating-point differences.

License and related releases

Research use is governed by CC BY-NC 4.0, with upstream attribution and the bundled DeBERTa MIT notice retained. See LICENSE, NOTICE and vendor/. Dataset licenses and terms apply separately.

• GRASP
• GRASP-A1
• GRASP source repository

The loader source for this release is inside grasp_molecular-0.1.1-py3-none-any.whl; install this wheel for A2 rather than an older source-repository revision.