GRASP-A1 / README.md

Add paired GRASP/A1 Colab comparison with shared audited splits

b405a6e verified 10 days ago

8.36 kB

---
license: cc-by-nc-4.0
library_name: pytorch
tags:
  - chemistry
  - molecular-representation-learning
  - graph-transformer
  - replaced-token-detection
  - assay-supervision
  - grasp
  - ablation
---
# 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](https://github.com/caithmac/GRASP). This repository uses its custom `GRASPEncoder` loader; a standard Transformers `AutoModel` is not provided.

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

The repository is public and requires no authentication:

```python
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:

```python
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.

```bash
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](https://colab.research.google.com/drive/15G8jCYj2RmRiJLUtjoZqi_VHvzw89GIF), or [download its self-contained source](notebooks/01_compare_GRASP_A1.ipynb) 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](notebooks/verification.json) 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](https://github.com/caithmac/GRASP).

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.

	---
	license: cc-by-nc-4.0
	library_name: pytorch
	tags:
	- chemistry
	- molecular-representation-learning
	- graph-transformer
	- replaced-token-detection
	- assay-supervision
	- grasp
	- ablation
	---

	# 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](https://github.com/caithmac/GRASP). This repository uses its custom `GRASPEncoder` loader; a standard Transformers `AutoModel` is not provided.

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

	The repository is public and requires no authentication:

	```python
	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:

	```python
	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.

	```bash
	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](https://colab.research.google.com/drive/15G8jCYj2RmRiJLUtjoZqi_VHvzw89GIF), or [download its self-contained source](notebooks/01_compare_GRASP_A1.ipynb) 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](notebooks/verification.json) 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](https://github.com/caithmac/GRASP).

	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.