ChemGPT-Base: Transformer for Molecular Design

• ChemGPT-Base is a transformer-based generative model that uses SMILES tokenization and masked self-attention to design valid molecules with controlled properties.
• It utilizes a lightweight, six-layer decoder with 512-dimensional embeddings and eight attention heads, optimized through autoregressive training and teacher-forcing.
• Benchmark results on the MOSES suite show high chemical validity (>98.5%), multi-property satisfaction (~78%), and efficient generation times, making it a strong baseline for molecular discovery.

ChemGPT-Base denotes a family of transformer-based generative models for molecular design conditioned on chemical language representations, particularly SMILES (Simplified Molecular Input Line Entry System). Developed in the wake of the Generative Chemical Transformer (GCT) paradigm, ChemGPT-Base leverages masked self-attention to learn chemical structure, enforce chemical validity, and support property-conditioned de novo molecule generation. The canonical configuration is a lightweight decoder-only transformer trained on large molecular corpora with property-conditioning. ChemGPT-Base is widely cited as the foundation for high-validity, property-controllable chemical language modeling and serves as a reference point in benchmarking generative models for molecular discovery (Kim et al., 2021).

1. Model Architecture and Tokenization

ChemGPT-Base adopts a stack of six Pre–Layer-Norm transformer decoder blocks, each with eight attention heads and 512-dimensional hidden states. Input molecules are encoded via SmilesPE, a substructure-driven tokenization yielding a 28-token vocabulary with atom, bond, branch, and ring tokens. Variable-length SMILES are padded to a maximum length of 80 tokens. Token embeddings and sinusoidal positional encodings are summed before entering the transformer layers. Each block includes multi-head masked self-attention, positionwise feed-forward layers with hidden size 2048, and dropout (rate 0.3). The attention mechanism is responsible for modeling both local and long-range token dependencies inherent to chemical graphs, despite SMILES’ inherent semantic discontinuities.

Mathematically, for each head $i$, query ($Q_i$), key ($K_i$), and value ($V_i$) projections are applied, with scaled dot-product attention:

$$\text{Attention}(Q_i,K_i,V_i) = \mathrm{Softmax}\left(\frac{Q_i K_i^{\top}}{\sqrt{d_k}}\right) V_i$$

The transformer output passes through a positionwise feed-forward network:

$$\operatorname{FFN}(H) = \max(0, H W_1 + b_1) W_2 + b_2$$

where $W_1, W_2$ are trainable weight matrices.

2. Training Objectives and Optimization

ChemGPT-Base is optimized as an autoregressive LLM using teacher-forcing. The model decomposes the joint molecular sequence probability:

$P(x) = \prod_{t=1}^T P(x_t \mid x_{<t}, c)$

where $x = (x_1,\ldots,x_T)$ are SMILES tokens and $c$ is the optional conditioning vector. The principal loss is the cross-entropy over tokens:

$Q_i$0

Regularization includes dropout (0.3 on all transformer layers) and weight decay ($Q_i$1). In variants inspired by GCT, a KL-divergence term is added for VAE-style latent code sampling, but this is omitted in the standard ChemGPT-Base.

Key hyperparameters for model training are summarized:

Hyperparameter	Value
Layers	6
Attention heads	8
Hidden dimension	512
Feedforward dimension	2048
Dropout	0.3
Batch size	256
Optimizer	Adam ($Q_i$2, $Q_i$3)
Learning rate	$Q_i$4, warm-up schedule
Epochs	25
Warm-up steps	100,000
Weight decay	$Q_i$5

Training typically completes in $Q_i$672 hours on 4$Q_i$7NVIDIA V100 GPUs (16 GB each) (Kim et al., 2021).

3. Property Conditioning and Controlled Generation

ChemGPT-Base supports explicit property conditioning for molecular inverse design. Properties—typically logP, topological polar surface area (tPSA), and QED—are computed via RDKit, scaled to $Q_i$8, and projected to 512-dimensional embeddings using a learnable MLP. Conditioning vectors are prepended as control tokens to the input or injected at each layer via concatenation, depending on the architectural variant.

Generation is performed using:

• Greedy decoding (argmax) for rapid inference,
• Beam search (beam size 4) for improved log-likelihood,
• Top-$Q_i$9 ($K_i$0) or nucleus sampling ($K_i$1) for diversity.

Explicit grammar and valence checks (via RDKit) follow every generation step, ensuring that generated SMILES are chemically valid. A valence-aware mask further restricts improbable next-token choices during sampling.

4. Evaluation Metrics and Benchmarking

ChemGPT-Base is evaluated using the MOSES benchmark suite on a set of 30,000 generated molecules. Metrics include validity, uniqueness, novelty, internal diversity (IntDiv), Fréchet ChemNet Distance (FCD), fragment/scaffold similarity, SNN, and property-matching rate (PMR):

Metric	ChemGPT-Base (GCT)	Best Reference Model
Validity	0.9853	0.9767 $K_i$2 0.0012 (VAE)
Unique@1k	1.0000	1.0000 (all)
Unique@10k	0.9981	0.9984 $K_i$3 0.0005 (VAE)
IntDiv	0.8531	0.8565 $K_i$4 0.0007 (AAE)
Novelty	0.8144	0.9498 $K_i$5 0.0006 (CharRNN)
PMR (all props)	78.3%	Not specified
MAE (logP)	0.177	-
MAE (tPSA)	2.923	-
MAE (QED)	0.035	-

Validity regularly exceeds 98.5%, uniqueness approaches 99.8%, and multi-property satisfaction is approximately 78%. Molecule generation times are 440–507 ms/sample on modern GPUs.

5. Self-Attention: Chemical Structure Learning

Self-attention in ChemGPT-Base captures both short- and long-range token interactions. Notably, attention heads can focus on atom pairs adjacent in molecular graphs yet separated by $K_i$610 tokens in SMILES. This reflects the mechanism’s ability to bridge SMILES “semantic gaps,” reconstructing implicit graph locality. Attention maps exhibit high mutual weight between distant but chemically adjacent tokens, raising the rate of valid generations above 98% despite SMILES’s non-local encoding. This structural learnability is a foundational advantage over recurrent or n-gram methods (Kim et al., 2021).

6. Applicability and Limitations

ChemGPT-Base, with its 6-layer, 512-dim, 8-head Pre-LN transformer backbone, achieves an effective trade-off between validity, property-guided control, and computational efficiency. The architecture is sufficiently lightweight for sub-second generation on commodity GPUs yet expressive enough for chemically diverse, valid, and property-conditioned de novo molecular design. These properties establish ChemGPT-Base as a practical baseline for research in language-driven chemical generation and as a precursor to larger multi-task frameworks such as ECNU-ChemGPT (Zhang et al., 2 Jul 2025).

A plausible implication is that the transformer attention structure, when combined with robust chemical tokenization and property conditioning, enables scalable extension to more complex tasks including retrosynthesis prediction, reaction outcome modeling, and chemical question answering, as further explored in subsequent chemistry-specialized LLMs.