ZeoBench: A Benchmark for Self-Supervised Learning on 3D Zeolite Representations

AI summary

Overview

Research area: 3D computer vision applied to materials science — self-supervised representation learning on volumetric representations of zeolite crystals.

Technical level: Intermediate. Readers should know what contrastive and distillation-based self-supervised learning are (SimCLR, MoCo, BYOL, DINO), and be comfortable with 3D convolutional networks; no chemistry background is required.

Scope: A benchmark paper that standardizes pretraining and evaluation of 3D self-supervised learning on all-silica zeolite distance grids, comparing four SSL objectives against out-of-domain 3D pretraining, ImageNet-pretrained 2D backbones adapted to 3D, hand-crafted descriptors, and supervised baselines on six adsorption prediction tasks.

What This Paper Is About

Zeolites are porous crystalline materials whose adsorption behavior is determined purely by 3D pore geometry, but property labels are expensive to obtain through simulation (roughly one hour per structure–molecule pair) and are only available for limited structure–molecule pairs. As a result, most of the 331k known real and hypothetical zeolite structures sit unused, while supervised models train on small labeled subsets and often require on the order of 5,000–10,000 labeled examples to work well. The paper asks which self-supervised pretraining strategies produce the best transferable 3D representations in this low-label regime, and packages the answer as a standardized benchmark.

Key Contributions

1. ZeoBench benchmark: a self-supervised learning benchmark on dense 3D grids of all-silica zeolites, built from 331k real and hypothetical zeolite structures (402 experimentally synthesized, 331,172 theoretically feasible) with fixed pretraining/validation/test splits of 301k / 20k / 10k.
2. Standardized downstream protocol: six existing adsorption-property prediction tasks with label-budget splits from 100 to 10,000 structures, defined metrics, and both linear-probe and fine-tuning protocols.
3. Broad baseline comparison: four SSL objectives (BYOL, DINOv1, MoCov2+, SimCLR), out-of-domain 3D volumetric pretraining (3DINO), ImageNet-pretrained 2D ConvNets and ViTs adapted to 3D via three strategies (2D slices, channel inflation, depth tokenization), hand-crafted descriptors, and from-scratch supervised baselines.
4. Augmentation and architecture analysis: ablations showing that scale augmentation is critical for SSL pretraining and that CNNs currently outperform ViTs on this task, plus a transfer experiment to metal organic frameworks.

Main Findings

• In-domain SSL sets a new state of the art: On the C18 full training set (70k), SSL-pretrained backbones beat the prior supervised ZeoNet baseline: MoCo 3.1 ± 0.2, DINO 3.2 ± 0.2, SimCLR 3.3 ± 0.2, BYOL 3.6 ± 0.2 versus ZeoNet 3.7 ± 0.2 (MSE of log(k_H), lower is better). The authors state the best pretraining method outperforms training from scratch on C18 by 16% on the full dataset.
• The gap widens as labels shrink: Pretrained models outperform from-scratch training more strongly at small subset sizes. At the 10k subset, fine-tuning average MSE is 11.8 for SimCLR, 12.4 for BYOL, 12.6 for MoCo, 13.4 for DINO and 13.5 for a from-scratch ResNet.
• Contrastive methods give the best frozen features: In linear probing, SimCLR (28.5 average) and MoCo (27.4) clearly beat DINO (44.1) and BYOL (34.3), which the authors attribute to the more uniform embedding distribution produced by explicit negatives — potentially better suited to continuous regression than to classification.
• The ranking shifts after fine-tuning: Despite the large linear-probe margin, BYOL becomes comparable to the contrastive methods once the backbone is fine-tuned, suggesting distillation methods learn useful features that are less directly usable when frozen.
• ImageNet is the strongest non-SSL baseline: Adapted ImageNet features (best variant: depth tokenization, 48.7 average linear-probe MSE) beat hand-crafted descriptors (56.9) and medical-volume 3DINO pretraining (60.8), though they require substantial test-time augmentation (averaging over 20 translations and rotations) to handle channel offsets and molecule orientation.
• Hand-crafted descriptors degrade in low-label regimes: Features such as pore-limiting diameter, largest-cavity diameter, surface area, pore volume and Si density are competitive at the 10k subset but overfit with few training points.
• Scale augmentation is essential: Pretraining with only translation and rotation (σ = 0) performs much worse than with scale augmentation on C18. The strength of the scale noise (σ varied from 0 to 0.4) has little effect; only its presence matters.
• Invariance trades off against accuracy: Networks pretrained with scaling are less geometrically invariant (measured as standard deviation of predictions across 100 augmented views) yet perform better downstream, indicating a tradeoff between strict equivariance and transferable representations.
• CNNs beat transformers: From-scratch Swin (38.4) and ViT (40.4) average MSE are far worse than from-scratch ResNet (13.5), and transformers show much lower translational invariance after training, possibly because the dataset is too small for them.
• Pretraining transfers out of domain: On the hMOF dataset, a ZeoBench SimCLR backbone fine-tuned with 5,000 training MOFs reaches 0.242 average MAE versus 0.375 from scratch, and with the full 96,356 training set reaches 0.121 versus 0.124 for scratch and 0.241 for the MOF-specific Barlow Twins MOFormer baseline.

Methodology in Plain English

The authors convert each zeolite into a single-channel "signed distance grid": a 3D voxel volume where every voxel stores the distance to the nearest solvent-accessible surface, with negative values inside. Base grids are computed at 0.15 Å resolution and then repeated and downsampled with trilinear interpolation to the target resolution (0.45 Å by default).

They take 301k structures for pretraining and hold out 20k validation and 10k test structures, never touched during pretraining. All four SSL methods use the same 3D ResNet-18 backbone and the same augmentations — random 3D translations, rotations, and scale changes implemented by jittering the grid resolution as Δ_d = max(0.45 + N(0, σ²), 0.05) Å — so that differences reflect the objective, not the setup. The two contrastive methods (SimCLR, MoCo) use the InfoNCE loss; the two distillation methods (BYOL, DINO) match a momentum-updated teacher without explicit negatives. Pretraining runs with LARS, batch size 1024 across 2 NVIDIA A100 GPUs for 400 epochs (about 3.5 days), built on the solo-learn library with 64×64×64 volumes.

Evaluation happens two ways. In linear probing, the backbone is frozen and only a linear head is trained on unaugmented grids (100 epochs). In fine-tuning, the whole network trains with augmentations (50 epochs, AdamW, single GPU). Downstream targets are Henry's constants for six hydrocarbons — three linear alkanes C18, C24, C30 and three mono-/di-branched C18 isomers (2C17, 4C17, 22C16) — with labels subsampled into nested subsets from 100 to 10,000 structures and evaluated as MSE of log(k_H) with bootstrapped error bars over 1000 trials.

To test whether 2D ImageNet features could substitute for in-domain 3D pretraining, they try three adaptations: taking three 2D slices as image channels, inflating the first convolutional layer's weights to accept D channels (downscaled by √(3/D)), and "depth tokenizing" the volume into 224×224×3 chunks fed sequentially to a ViT. They then repeat the pipeline on metal organic frameworks, building van der Waals-radius distance grids for the hMOF dataset.

Why This Matters

Impact on research: The paper provides the first standardized testbed for self-supervised learning on dense 3D molecular grids of porous materials, with fixed splits, released pretraining configurations and code, so that new SSL objectives, architectures and augmentations can be compared fairly. It also provides a counterpoint to the assumption that in-domain pretraining always wins, since ImageNet features are the strongest non-SSL baseline here.

Real-world applications:

• Gas storage and separation: identifying zeolites that selectively adsorb CO2 or CH4 at relevant partial pressures, which the hMOF transfer experiment targets.
• Catalysis: screening frameworks for pore geometries compatible with a target molecule before committing to synthesis.
• Materials discovery under scarce labels: prioritizing which of the 331k candidate structures are worth expensive simulation or experimental characterization.
• Cross-material screening: reusing a zeolite-pretrained backbone on metal organic frameworks, where chemical composition differs but the geometric reasoning transfers.

Industry relevance: Chemical and petrochemical companies rely on adsorbents and catalysts; a pretrained backbone that predicts adsorption from geometry with much lower runtime than Monte Carlo simulation, and that works with only hundreds of labels, shortens the screening loop for new framework candidates.

Future Directions

• Scale beyond a fixed backbone: the study holds the 3D ResNet-18 fixed to isolate objective effects; the authors explicitly call for scaling studies over larger backbones, longer pretraining, larger batch sizes and more compute, and note that parameter–performance scaling laws have yet to be shown on ZeoBench.
• Make transformers competitive: ViTs and Swin lag well behind CNNs, which the authors attribute partly to dataset size; transformer-specific reconstruction objectives such as masked autoencoding are untested.
• Group-equivariant networks: architectures that build in rotation and translation symmetry are appealing for 3D grids but hard to scale, since activation memory and compute grow roughly linearly with group size K.
• Broaden tasks and joint structure–molecule modeling: expanding to other frameworks, target molecules and downstream tasks (catalysis, separation), and encoding porous structures and target molecules jointly across zeolites and MOFs.

Target Audience

Graduate students and researchers working at the intersection of 3D computer vision and scientific machine learning, particularly those interested in self-supervised learning for scientific volumetric data, geometric deep learning on crystals, or low-label regression transfer. It is also useful for computational chemists and materials scientists who want to know which representation-learning approach to adopt for adsorption screening, and for practitioners deciding between CNNs and transformers on small 3D scientific datasets.

Authors’ abstract

Zeolites are an important class of porous crystalline materials with broad applications in gas storage, separations, and catalysis. Although hundreds of thousands of experimentally realized and hypothetical zeolite structures are known, identifying materials optimized for a target application remains challenging because property labels are expensive to obtain and available only for limited structure-molecule pairs. While prior work has demonstrated that 3D neural…