BRAIDBERTa / README.md
aakothari's picture
Update README.md
104990c verified
|
Raw
History Blame Contribute Delete
7.56 kB
---
license: cc-by-4.0
---
# BRAIDBERTa
BRAIDBERTa is the first pretrained language model for BRAID, a molecular notation in which every valid string decodes to a valid molecule.
Compared against a matched SMILES baseline on BBBP, BRAID achieves statistically indistinguishable ROC-AUC while guaranteeing valid decoding.
Check it out: https://huggingface.co/spaces/aakothari/Conformer
---
## The BRAID notation
BRAID (**B**ranch-counted, **R**elative-ring, **A**ttachment-native, **I**nvalid-free,
**D**ense) is built from four mechanisms, all of which are recombined prior art:
| Mechanism | Syntax | Borrowed from |
|---|---|---|
| Organic-subset lexing, `[...]` only when needed | `C`, `Cl`, `[N+]` | SMILES |
| Length-counted branches (no parentheses) | `>k` then *k* atoms | SELFIES branch counts |
| Relative ring closure (no paired digits) | `^d` = bond *d* atoms back | DeepSMILES relative rings |
| Fragment separator, global index continues across it | `.` | SAFE |
Only **atom tokens** advance the global index; `>k`, `^d`, bond symbols and `.` do not.
```
aspirin SMILES CC(=O)Oc1ccccc1C(=O)O
BRAID CC>1=OOcccccc^5C>1=OO (aromatic mode)
BRAID CC>1=OOC=CC=CC=C^5C>1=OO (kekulé mode, default)
```
### Novelty: none claimed
BRAID is **not novel at the mechanism level.** Relative ring closure is DeepSMILES (O'Boyle
& Dalke, 2018). Length-counted branches are the SELFIES branch-count idea (Krenn et al.,
2020) with the brackets stripped. Guaranteed-valid decoding via a valence state machine is
the entire point of SELFIES. The fragment separator is SAFE (Noutahi et al., 2023).
Please treat BRAID as a recombination.
### The property BRAID does have
Every string decodes. On 10,000 random and mutated strings, the decoder produced a
sanitizable molecule with **0 crashes**. This is achieved by making the decoder *total*:
unparseable tokens and over-valent bonds become silent no-ops. Validity is guaranteed.
Under an order-*k* Markov generator on a 79-molecule corpus:
| representation | valid% | unique% | novel% | vocab | tok/mol |
|---|--:|--:|--:|--:|--:|
| SMILES | 40.0 | 38.4 | 34.0 | 24 | 13.6 |
| DeepSMILES | 35.4 | 46.9 | 42.4 | 27 | 13.3 |
| SELFIES | 100.0 | 59.6 | 57.9 | 29 | 12.9 |
| **BRAID** | **100.0** | 60.9 | 59.2 | 30 | 13.2 |
A Markov model is far weaker than a real CLM, so treat validity% as the meaningful
(model-independent) signal and uniqueness/novelty as illustrative.
---
## Model details
- **Architecture:** RoBERTa (masked language modelling)
- **Pretraining corpus:** ZINC 100k, encoded in BRAID
- **Tokenizer:** BPE trained on the BRAID corpus
- **Objective:** MLM
---
## Downstream evaluation: BBBP
Finetuned for binary classification (blood–brain barrier penetration) with a fresh 2-logit
head and cross-entropy loss. Hyperparameters (learning rate, epochs) were selected by Optuna
on the validation split, then the best configuration was retrained across 5 seeds. Metric is
test ROC-AUC.
| Model | Notation | BBBP ROC-AUC | Seed 0 | Seed 1 | Seed 2 | Seed 3 | Seed 4 |
|---|---|---|---|---|---|---|---|
| **BRAIDBERTa-v9** | BRAID | **0.725 ± 0.017** | 0.727 | 0.710 | 0.751 | 0.710 | 0.727 |
| *SMILES control* | SMILES | *0.720 ± 0.010* | 0.719 | 0.738 | 0.715 | 0.714 | 0.712 |
(± is the sample standard deviation across seeds.)
### Interpretation
**The two are statistically indistinguishable.** Welch's *t*-test on the five seeds per
model gives *t* ≈ 0.61, *p* ≈ 0.56. The 0.005 gap in favour of BRAID sits well inside seed
noise, and the per-seed ranges overlap almost completely (BRAID 0.710–0.751, control
0.712–0.738).
The correct reading is
> Under matched pretraining and finetuning, BRAID's guaranteed-validity property comes
> **at no measurable cost to downstream predictive performance** on BBBP.
Note: The seed-to-seed spread (±0.017) exceeds the between-model difference. Any comparison
at this scale reporting a single seed is measuring noise.
---
## Usage
The model consumes **BRAID strings, not SMILES**, so you need the codec:
```bash
pip install rdkit transformers
pip install git+https://github.com/AayushK-othari/braid.git
```
```python
from braids import smiles_to_braid
from transformers import AutoTokenizer, AutoModelForSequenceClassification
tok = AutoTokenizer.from_pretrained("aakothari/BRAIDBERTa-v9")
model = AutoModelForSequenceClassification.from_pretrained(
"aakothari/BRAIDBERTa-v9", num_labels=2
)
braid = smiles_to_braid("CC(=O)Oc1ccccc1C(=O)O") # -> CC>1=OOC=CC=CC=C^5C>1=OO
inputs = tok(braid, return_tensors="pt", truncation=True, max_length=128)
logits = model(**inputs).logits
```
> **Do not feed raw SMILES to this tokenizer.** It will not error — BPE always backs off to
> characters — it will simply produce a meaningless tokenization and a plausible-looking,
> wrong prediction. Always encode with `smiles_to_braid` first.
**Encode in the same mode the model was pretrained in.** The codec has Kekulé (default) and
aromatic modes, and they produce different strings for the same molecule
(`C=CC=CC=C^5` vs `cccccc^5`). Mixing modes between pretraining and inference silently
degrades performance.
<!-- TODO: state which mode the ZINC 100k pretraining corpus used -->
The reference implementation — encoder, decoder, valence state machine, stereo handling,
tokenizer/`Vocab` builder, and the test suites — lives at
**[github.com/aakothari/braids](https://github.com/aakothari/braids)**.
### Verified codec behaviour
- **Round-trip:** 42/42 molecules recover their constitution in both Kekulé and aromatic
modes — including fused/bridged systems (naphthalene, indole, adamantane, caffeine,
purine, quinoline, nicotine) and ions (nitromethane, acetate, ammonium, Na⁺·Cl⁻).
- **Stereo:** 26/26 chiral / E-Z molecules round-trip in both modes — L-alanine,
(S)-ibuprofen, meso vs. L-tartaric acid, menthol, L-DOPA, adrenaline, carvone,
fumaric/maleic acid, conjugated dienes, mixed chiral + E/Z.
- **Validity:** 10,000/10,000 random and mutated strings decode to a sanitizable molecule
with 0 crashes.
---
## Limitations
1. **Stereochemistry is partial.** Tetrahedral chirality and E/Z double bonds round-trip on
26/26 test cases. Allene/axial/planar chirality, atropisomers, and non-tetrahedral stereo
centres fall back to unspecified.
2. **The valence state machine is approximate.** A hand-rolled charge→valence rule, correct
for common ions but liable to mis-clamp hypervalent, organometallic, or unusual-charge
atoms. It is not a substitute for RDKit's model.
3. **Insertion/deletion is non-local.** Inserting or deleting an atom shifts every `^d` that
spans the edit point. Inherent to relative back-references; applies equally to DeepSMILES.
4. **Not a canonical hash.** One molecule has many valid BRAID strings. The encoder uses
RDKit canonical-rank rooting for determinism, but uniqueness is not proven.
---
## Citation
BRAID is a recombination of published mechanisms. Please cite the underlying work:
- O'Boyle, N. & Dalke, A. *DeepSMILES: An Adaptation of SMILES for Use in Machine-Learning
of Chemical Structures.* ChemRxiv, 2018.
- Krenn, M. et al. *Self-Referencing Embedded Strings (SELFIES): A 100% robust molecular
string representation.* Machine Learning: Science and Technology, 2020.
- Noutahi, E. et al. *Gotta be SAFE: A New Framework for Molecular Design.* 2023.
- Weininger, D. *SMILES, a chemical language and information system.* J. Chem. Inf. Comput.
Sci., 1988.