GRASP-A1

GRASP-A1 is the GRASP molecular graph Transformer trained without absolute atom-position embeddings. It retains shortest-path relative attention, radius-0 atom-environment tokens, and the replaced-token-detection objective. The encoder has 93,091,584 parameters, 12 layers, hidden size 768, and 12 attention heads.

This research release provides the completed 50,000-update Step2 encoder by default and the corresponding Step1 encoder under step1/. Step1 saw 1,538,242,560 molecule presentations from ZINC20; this is a presentation budget, not a count of certified unique molecules. Step2 used the audited 511,897-molecule ChEMBL 36 set with 642 assay tasks.

Use

Install Python 3.10โ€“3.13 and the GRASP code repository. This repository uses its custom GRASPEncoder loader; a standard Transformers AutoModel is not provided.

python -m pip install "git+https://github.com/caithmac/GRASP.git@7843e72606dd208517cd8cf86ee07b3bb965cc79"

The repository is public and requires no authentication:

from grasp import GRASPEncoder

encoder = GRASPEncoder.from_pretrained("caithmac/GRASP-A1")
vectors = encoder.encode(["CCO", "c1ccccc1"])
print(vectors.shape)  # (2, 768)

The output is the final-layer CLS state. A property prediction requires a fitted endpoint head. The default embedding is different from the learned layer-mixture and atom-attention readout used in the primary OpenADMET evaluation.

To test Step1 before assay adaptation:

from pathlib import Path
from huggingface_hub import snapshot_download
from grasp import GRASPEncoder

snapshot = snapshot_download("caithmac/GRASP-A1", allow_patterns="step1/*")
encoder_s1 = GRASPEncoder.from_pretrained(Path(snapshot) / "step1")
vectors_s1 = encoder_s1.encode(["CCO", "c1ccccc1"])

Fine-tune on your tasks

Create separate train.csv and valid.csv files with smiles and target columns. Keep test molecules separate from training and model selection. Regression targets use their original units; binary targets must be 0 or 1.

python -m grasp finetune --model caithmac/GRASP-A1 \
  --train-csv train.csv --valid-csv valid.csv \
  --smiles-column smiles --target-column target \
  --task regression --method full --output-dir runs/a1_my_task

python -m grasp predict --model runs/a1_my_task \
  --input-csv test.csv --output-csv predictions.csv

Use --task binary for binary classification. --method frozen trains the readout and head with a fixed encoder; --method lora provides LoRA adaptation. Use a separate output directory per endpoint. To compare against published GRASP, repeat the same split, preprocessing, method and seed with --model caithmac/GRASP.

Compare A1 and original GRASP in Colab

Open the comparison notebook in Colab, or download its self-contained source and use Colab โ†’ File โ†’ Upload notebook. Select a GPU and run cells in order. Public access requires no token; if access is restricted later, the notebook supports the Colab secret HF_TOKEN or a hidden login prompt.

The notebook tests ESOL, FreeSolv and BBBP with shared scaffold splits, matched initial heads and training settings, and pinned checkpoint revisions. It exports paired metrics, per-molecule predictions, split hashes and data-exclusion audits. Defaults use one training seed; set SEEDS = [0, 1, 2] for repeated training on the same splits. This compares the released checkpoints and does not isolate the positional-embedding change or reproduce paper scores. See notebook validation for the real-data split audits and tiny CPU workflow checks; full Colab GPU benchmarks have not been run as part of notebook creation.

Model details and provenance

Property Value
Architecture 12 layers, width 768, 12 heads, FFN width 3072
Encoder parameters 93,091,584
Vocabulary 211 radius-0 atom-environment and special tokens
Absolute position embeddings Removed throughout training
Relative attention Graph shortest-path, content-to-position (c2p)
Step1 objective Replaced-token detection, 25% selected atom positions, GDES
Step1 budget 1,538,242,560 molecule presentations
Step2 50,000 updates; sparse supervision on 642 ChEMBL 36 assays
Default representation Final-layer CLS, 768 dimensions

model.safetensors contains Step2 encoder tensors; step1/model.safetensors contains the effective GDES discriminator encoder before Step2. Generator weights, optimizer states, and fitted downstream heads are omitted. Each SafeTensors file is checked tensor by tensor against its encoder export. training_provenance.json records source checkpoint hashes and architecture; SHA256SUMS records release file hashes. validation_report.json describes the actual loading and adaptation smoke checks, which are not benchmark evaluations.

Evaluation

The matched control/A1 experiment completed 23 OpenADMET endpoints with three split/seed runs per endpoint, using endpoint-specific full fine-tuning with a learned layer mixture and atom-attention pooler:

Model in the matched experiment Mean endpoint MAE Endpoints with lower mean MAE
Retrained control with absolute positions 0.370232 12/23
A1 without absolute positions 0.368588 11/23

The MAE difference is โˆ’0.001644, about 0.44% relative to the matched control. This is mixed, near-parity evidence rather than a general superiority claim. These results come from one pretraining seed; downstream split and seed change together. The matched control is a separately retrained comparator, and these numbers are not a direct benchmark against the public caithmac/GRASP weights or zero-shot encoder scores. TDC evaluation and fixed-split training repeats were still incomplete when this release was prepared. AqSolDB remains blocked because two official test molecules failed parsing; its test set was not modified.

Removing absolute positions addresses atom-order dependence. Existing direct atom-permutation audits measured relative embedding drift on the order of 1e-6 for A1. Numerically equivalent embeddings can still differ slightly with batching and floating-point arithmetic. Chemical standardization or parsing changes can change the graph and are outside this invariance statement.

Intended use and limitations

Use this model for research on molecular embeddings and property prediction after endpoint-specific fitting. The vocabulary and graph distances still omit tetrahedral and E/Z assignments: A1 remains stereo-blind. It does not supply explicit bond-type inputs, 3D conformers, protonation-state ensembles, or measured uncertainty. Improvements proposed for A2/A3 are not part of this checkpoint.

Structural pretraining used molecules with at most 96 heavy atoms. The GRASP loader accepts at most 511 atoms and rejects larger or invalid inputs explicitly; accepting a molecule does not establish that it is in the training domain. Standardization, tautomer handling, charged fragments and salts should be considered consistently when defining a task. This encoder is not a calibrated assay predictor and does not provide clinical or toxicological determinations.

Training data, license and citation

Structural pretraining used ZINC20; assay adaptation used ChEMBL 36. Raw datasets are not redistributed. Model and MolE-derived implementation terms are CC BY-NC 4.0; bundled DeBERTa code retains its separate MIT license. See LICENSE and NOTICE and the GRASP source and reproduction records.

Satya Pratik Srivastava, Rohan Gorantla, Sharath Krishna Chundru, Harshit Singh, Antonia S. J. S. Mey, and Rajeev Kumar Singh. GRASP: Graph Representation Learning with Assay Supervision for Molecular Properties. Use the citation metadata in the GRASP code repository. A1 is an experimental ablation checkpoint, not a separately established paper result.

Downloads last month
18
Safetensors
Model size
93.1M params
Tensor type
F32
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support