Sphere-Native Transformers for Molecular Property Prediction

When you train a chemical language model on SMILES strings the way you'd train one on English, you are asking the network to rediscover chemistry from scratch. You give it the alphabet, C, O, N, =, (, ), and you give it tens of millions of molecules, and you hope that, somewhere in the gradient signal, the model figures out valences, bond orders, conjugated rings, and multipole structure on its own.

It usually does. That's why ChemBERTa, MolFormer, and their successors work. But it works the way scaling laws always work: by spending compute, data, and parameters on a problem that has a much more compact mathematical structure underneath.

At MSTACK AI we are exploring alternative paradigms. We are an AI-native chemistry R&D company. Every layer of our product, from data ingestion to property prediction to active-learning loops in our flagship ChemStack AI platform, is built on the conviction that chemistry's own mathematical priors should be in the architecture, not learned around the architecture. The result is models that are smaller, faster, cheaper to train, and more accurate on the endpoints our customers care about.

Our first foundational research paper, Chem-GMNet, is the proof point. It is now on arXiv. This post explains what it does, why it works, and what it means for chemistry-AI deployments at industrial scale.

The picture in one paragraph

Chem-GMNet is a transformer family in which every standard module, token embedding, attention, feed-forward layer, output head, is replaced by a counterpart that operates on the unit hypersphere S^(k−1). Each SMILES token is a learnable direction on the sphere. Each layer mixes those directions through a Gegenbauer-kernel attention whose persistent state is, provably, the truncated multipole expansion of the molecular distribution. Each feed-forward layer is a spherical convolution whose nonlinearity has been compiled, at initialisation, into per-harmonic eigenvalues. The whole architecture is built on three classical results, Stone-Weierstrass, Schoenberg's theorem on positive-definite zonal kernels, and the multipole theorem, and each of them does concrete work inside the model.

Why chemistry deserves its own geometry

Most modern chemical AI takes one of two routes. The dominant route treats SMILES strings as ordinary text and trains very large transformers, ChemBERTa-2, MolFormer, MoLFormer-XL, ChemBERTa-3, with self-supervised objectives on tens of millions of molecules. This works because of scale. The alternative route exploits chemistry's geometry directly through 3D-equivariant graph neural networks, SchNet, DimeNet, NequIP, MACE, Equiformer, but those models require precomputed 3D conformers as input, which are expensive to generate and not always available.

Both approaches leave money on the table.

Generic transformers ignore the fact that chemistry's entities live on natural manifolds. A molecular orbital is a function on a sphere. A charge distribution decomposes cleanly into multipole moments. A bond between two atoms has a well-defined angular character. The Schoenberg-Bochner machinery that physicists have used since 1942 to characterise positive-definite kernels on spheres is exactly the right mathematical object for "two molecules are similar if their angular features overlap." None of this structure is given to a SMILES-based language model. It is left to gradient descent to discover, slowly and expensively.

3D-equivariant message passing uses geometry but only the external geometry, the 3D positions of atoms. It says nothing about the internal geometry of the model's representation. The features inside the network still live in unconstrained Euclidean space.

Chem-GMNet is the third route. The features inside the model live on the unit sphere S^(k−1) from layer 1 to the output head. No external conformer requirement, no scaling sledgehammer.

The three classical theorems we put to work

Every architectural choice in Chem-GMNet is anchored to a piece of harmonic analysis that has been known for decades. Naming them up front is the discipline that keeps the design honest.

These three results have lived in textbooks for fifty to a hundred years. They have not, to our knowledge, been used as the design principle for a chemical language model. Chem-GMNet shows what happens when you do.

The three modules

Three sphere-native modules implement the architecture end-to-end.

SH-Embedding: tokens as learnable directions

The standard V × d embedding table is replaced by a learned per-token unit direction on S^(k−1) plus a small residual. The token's embedding is the spherical-harmonic feature map of that direction projected back to the model dimension. At chemistry vocabulary sizes, this replaces a dense lookup table with a structured projection through a 156-dimensional harmonic basis.

DualSKA: sphere-flow plus sphere-kernel, fused

The attention block runs two branches in parallel over identical Gegenbauer features and fuses them with a per-head learned gate. The first branch, gated Sphere Flow Attention (SFA), is linear-time in sequence length; its terminal state is provably the multipole expansion of the input distribution. The second, Sphere Kernel Attention (SKA), is a softmax over the same Schoenberg-valid Gegenbauer kernel and returns a renormalised aggregate direction on the sphere. The fusion gate is one parameter per head; the model learns whether each head wants linear long-range memory, sharp full-context attention, or a mix.

SH-FFN: a spherical convolution, Funk-Hecke compiled

Where a standard transformer's feed-forward layer is a wide nonlinear map (d → 4d → d through a GELU), SH-FFN is a Funk-Hecke convolution. At module initialisation we compile the chosen zonal activation (default GELU) into per-harmonic Gegenbauer eigenvalues via Gauss-Legendre quadrature. The runtime forward path is then one elementwise scale per harmonic, not a 4× wider Euclidean GELU. This is where Chem-GMNet's per-layer encoder block becomes dramatically lighter than a vanilla transformer block: roughly 20× fewer FFN parameters at matched architectural shape, the largest single source of Chem-GMNet's parameter savings.

The three modules share substantial machinery: the same Gegenbauer feature map, the same Schoenberg-PD guarantee, the same sphere projection. The architecture is coherent, not a collection of tricks.

What the numbers actually show

The results in the paper are split across two evaluation regimes, both on canonical MoleculeNet scaffold splits under the chemberta3-faithful protocol: same data, same splits, same fine-tuning recipe ChemBERTa-3 published.

Random-initialised, scratch-vs-scratch

Against same-shape ChemBERTa-2 architecture trained without any pretraining, Chem-GMNet wins on 7 of 10 MoleculeNet endpoints (BACE, BBBP, HIV, Tox21, SIDER, ClinTox, ESOL, FreeSolv, Lipophilicity, MUV) at approximately 35% fewer parameters. The margin varies by endpoint: substantial on the tasks where the model's geometric structure compounds with task structure, more modest elsewhere.

Pretrained, fair-corpus comparison

When we pretrain Chem-GMNet on the same 10M-SMILES corpus that ChemBERTa-2 10M-MLM used, our random-initialised lead persists into the pretrained regime: we match or beat the public ChemBERTa-2 10M-MLM release on 5 of 7 shared endpoints.

The (k, L) scaling result

The most striking single number in the paper is this: at k=10, L=3, random-initialised Chem-GMNet hits an ESOL RMSE of 0.938, below the pretrained ChemBERTa-2 MLM-10M number on the same endpoint. We are beating a model that saw 10 million molecules of pretraining, using an architecture that saw zero. The compactness story is real, but so is the quality story: chemistry-faithful geometry is not just a parameter-saver, it is an inductive bias that beats scale at the small-data limit where most industrial chemistry actually lives.

What is GM-Net, and what comes next

Chem-GMNet is one instantiation of a more general framework we call GM-Net, Geometric Measure Network. The name encodes the two commitments of the design. Geometric: the architecture lives on the hypersphere, not in unconstrained Euclidean space. Measure-theoretic: each input position is treated as a discrete signed measure on the sphere, and the architecture's persistent state is identifiable with the harmonic moments of that measure. The framework is a recipe; Chem-GMNet is the recipe applied to chemistry.

We are now extending GM-Net: LM-GMNet instantiates the framework for natural language modelling. We hope to put it up soon on the MSTACK AI Research blog and in upcoming arXiv preprints.

The deeper claim is that Schoenberg's theorem, Funk-Hecke's diagonalisation, and the multipole expansion are not chemistry-specific tools. They are the right mathematical machinery for any domain in which sequence elements carry meaningful geometric structure. Chemistry is the cleanest demonstration because the geometry is most explicit. Once the framework matures across domains, we will publish it as a unified result.

Try the paper

The full Chem-GMNet paper, including all proofs, experimental details, appendices, and the (k, L) ablation tables, is on arXiv. The codebase is being prepared for public release. If you are working in chemistry-AI and want to collaborate, reach out to the team.

If you are interested in the ChemStack AI B2B product, we are running a closed beta with selected pharma and materials partners through Q3 2026.

Chemistry has its own mathematics. It always did. We just got around to writing transformers that respect it.