Contrastive Language-Crystal Pretraining

Yang Jeong Park, Mayank Kumaran, Chia-Wei Hsu, Elsa Olivetti, Ju Li
Massachusetts Institute of Technology and University of Illinois Urbana-Champaign

This is the official implementation of Contrastive Language–Crystal Pretraining (CLaC). In this work, we introduce a contrastive pretraining framework for multimodal representation learning between materials science knowledge and crystal structures. CLaC aligns 3D crystal graphs with natural language descriptions by jointly training a crystal graph neural network and a text encoder in a shared embedding space. To overcome the scarcity and bias of literature-derived text for crystalline materials, the model is pretrained on property-conditioned synthetic narratives generated by large language models and grounded in computed materials properties. Through CLIP-style contrastive learning with both inter-modal and intra-modal objectives, CLaC learns semantically structured representations that enable zero-shot multimodal understanding and language-guided crystal retrieval and inverse design, without task-specific fine-tuning.

Repository Structure

CLaC-revision/
├── asset/                    # Images and architecture diagrams
├── config/                   # Configuration files (model, training, data)
├── data/                     # Data loading and augmentation modules
├── model/                    # Model definitions (CLaC, graph encoders, text encoders)
├── evaluation/               # Evaluation and analysis scripts
│   ├── compute_metrics.py              # Main evaluation (zero-shot, inverse design)
│   ├── compute_metrics_baseline.py     # LLM baseline evaluation
│   ├── analyze_text_embedding_similarity.py  # Text embedding analysis
│   ├── visualize_embeddings.py         # UMAP visualization
│   └── plot_*.py                       # Various plotting utilities
├── scripts/                  # Executable bash scripts
│   ├── evaluation/                     # Evaluation scripts
│   ├── visualization/                  # Visualization scripts
│   └── training/                       # Training scripts
├── tools/                    # Utility scripts and data preparation
│   ├── run_manual_sweep.py
│   └── query_materials.py
├── docs/                     # Documentation
│   ├── guides/                         # Usage guides
│   └── evaluation/                     # Evaluation methodology
├── outputs/                  # All experimental results
│   ├── figures/                        # Publication-ready figures
│   ├── component_comparison/
│   ├── inverse_design_dual/
│   ├── text_embedding_similarity/
│   └── umap_visualizations/
├── train.py                  # Main training script
├── sweep.yaml               # Hyperparameter sweep configuration
└── readme.md                # This file

System Requirements

Software Dependencies

• Operating System: Linux (tested on Ubuntu 20.04, 22.04), macOS (experimental support)
• Python: 3.10, 3.11 (tested on 3.10.12, 3.11.5)
• CUDA: 12.1 or 12.6 (tested on 12.1, 12.6)
• PyTorch: 2.8.0
• PyTorch Geometric: 2.6.1
• Core Dependencies:
  • orb-models==0.5.5 (Crystal graph neural network encoder)
  • jarvis-tools (Materials science utilities)
  • transformers (Hugging Face transformers for text encoding)
  • lightning (PyTorch Lightning for training)
  • hydra-core (Configuration management)
  • wandb (Experiment tracking)
  • spacy (Text processing, requires en_core_web_sm model)
  • umap-learn (Dimensionality reduction for visualization)
  • matplotlib, scikit-learn, scipy, pandas, python-dotenv

Hardware Requirements

• GPU: NVIDIA GPU with 16GB+ VRAM recommended (tested on NVIDIA A100 40GB, RTX 3090 24GB, V100 32GB)
  • Minimum: NVIDIA GPU with 12GB VRAM (training with smaller batch sizes)
  • For inference/evaluation only: 8GB VRAM sufficient
• RAM: 32GB+ system RAM recommended (minimum 16GB)
• Storage: 50GB+ free disk space (for dataset, model checkpoints, and outputs)
• CPU: Multi-core processor (8+ cores recommended for data loading)

Tested Configurations

The software has been successfully tested on:

• Ubuntu 22.04 LTS + CUDA 12.6 + PyTorch 2.8.0 + NVIDIA A6000 (48GB)
• Ubuntu 22.04 LTS + CUDA 12.1 + PyTorch 2.8.0 + NVIDIA RTX 3090 (24GB)

Installation

Prerequisites

• Python 3.10+
• CUDA 12.6 (or adjust installation commands for your CUDA version)
• Conda or Miniconda

Method 1: Step-by-Step Installation (Recommended)

1. Create and activate conda environment:

conda create -n clac python=3.10
conda activate clac

2. Install PyTorch with CUDA support:

# For CUDA 12.6 (recommended)
pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu126

# For CUDA 12.1
pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu121

3. Install PyTorch Geometric and extensions:

# For CUDA 12.6
pip install torch-geometric==2.6.1
pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv \
    -f https://data.pyg.org/whl/torch-2.8.0+cu126.html

# For CUDA 12.1, use: torch-2.8.0+cu121.html

4. Install all remaining dependencies:

pip install -r requirements.txt

5. Download spaCy language model:

python -m spacy download en_core_web_sm

Method 2: Using environment.yml (Alternative)

conda env create -f environment.yml
conda activate clac

# Install PyTorch with CUDA support
pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu126

# Install PyTorch Geometric and extensions
pip install torch-geometric==2.6.1
pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv \
    -f https://data.pyg.org/whl/torch-2.8.0+cu126.html

# Download spaCy language model
python -m spacy download en_core_web_sm

Note: Both PyTorch and PyTorch Geometric require CUDA-specific wheels, so they must be installed separately with the correct CUDA version. The requirements.txt and environment.yml include all other dependencies.

Verify Installation

Test your installation:

python -c "import torch; import torch_geometric; import orb_models; print('✓ Installation successful')"

Troubleshooting

CUDA version mismatch:

• Check your CUDA version: nvcc --version or nvidia-smi
• Install matching PyTorch version from PyTorch website

PyTorch Geometric installation issues:

• Ensure PyTorch is installed first
• Use the correct wheel URL for your CUDA version
• See PyG installation guide

Package conflicts:

• If you encounter dependency conflicts, create a fresh conda environment
• Ensure you follow the installation order: PyTorch → PyTorch Geometric → other packages

Installation Time

• Method 1 (Step-by-step): Approximately 15-30 minutes on a standard desktop computer with good internet connection
  • Conda environment creation: 1-2 minutes
  • PyTorch installation: 5-10 minutes (depending on internet speed)
  • PyTorch Geometric and extensions: 5-10 minutes
  • Remaining dependencies: 5-10 minutes
• Method 2 (environment.yml): Approximately 20-40 minutes (conda resolves all dependencies automatically)

Note: Installation time may vary depending on internet connection speed and system specifications.

Dataset

Data configuration

We have proposed the following data configurations for training CLaC:

1. Academic paper dataset: Text extracted from real research papers using materials as queries.
2. Synthesized Text: Text generated by leveraging the inference capabilities of LLMs based on material property values extracted from open materials databases.

We appreciate your understanding that we are unable to fully share all data and the corresponding pretrained model weights due to potential copyright concerns regarding academic literature. However, the synthesized data is publicly available at GPT-Narratives-for-Materials.

Download Dataset

Download the synthesized GPT narratives dataset:

# Create data directory
mkdir -p datafiles

# Download dataset from HuggingFace
wget https://huggingface.co/datasets/yjeong/GPT-Narratives-for-Materials/resolve/main/mp_3d_2020_gpt_narratives.parquet \
    -O datafiles/mp_3d_2020_gpt_narratives.parquet

# Alternative: using curl
curl -L https://huggingface.co/datasets/yjeong/GPT-Narratives-for-Materials/resolve/main/mp_3d_2020_gpt_narratives.parquet \
    -o datafiles/mp_3d_2020_gpt_narratives.parquet

Split into train/val/test sets:

python data/get_splits.py --input datafiles/mp_3d_2020_gpt_narratives.parquet --train-ratio 0.8 --val-ratio 0.1

This creates:

• datafiles/mp_3d_2020_gpt_narratives_train.parquet (80%)
• datafiles/mp_3d_2020_gpt_narratives_val.parquet (10%)
• datafiles/mp_3d_2020_gpt_narratives_test.parquet (10%)

Pre-training

To train the CLaC model, specify the base dataset path (e.g., datafiles/mp_3d_2020_gpt_narratives) in the configuration. The training code will automatically recognize and use the split files (_train.parquet, _val.parquet, _test.parquet) for training, validation, and testing.

Set configuration of training

See yaml files in config section. If you want to change training configuration, modify them before running.

Run pretraining

Training code is based on pytorch lightning and hydra configuration. To pretrain the CLaC model, run python train.py with appropriate configuration. On 4x NVIDIA A6000 GPUs, training for 5 epochs takes approximately 24 hours.

Advanced Features

Resume Training from Checkpoint

To continue training from a saved checkpoint:

# config/config.yaml
resume_from_checkpoint: outputs/experiment/epoch=49-step=5000.ckpt

Or via command line:

python train.py resume_from_checkpoint=outputs/experiment/epoch=49-step=5000.ckpt

See docs/guides/how_to_resume_training.md for details.

Evaluation for zero-shot tasks

In this work, we measured zero-shot ability of our CLaC. For zero-shot retrieval, we measured the model's performance in finding corresponding crystals when given texts, and vice versa. For zero-shot multimodal understanding, we utilized multiple-choice questions (included in GPT-Narratives-for-Materials) regarding crystals, such as their composition and structure.

Running Evaluations

Zero-shot retrieval and QA:

python evaluation/compute_metrics.py \
  --evaluation-method 'zero-shot retrieval' \
  --label text \
  --model-ckpt outputs/your-model/checkpoint.ckpt \
  --device cuda:0

Inverse design:

python evaluation/compute_metrics.py \
  --evaluation-method inverse-design \
  --property "band gap" \
  --k-values 3 10 100 \
  --model-ckpt outputs/your-model/checkpoint.ckpt \
  --device cuda:0

Batch evaluation scripts:

# Run all zero-shot evaluations
bash scripts/evaluation/zero_shot_qa_evaluation.sh

# Run inverse design evaluation
bash scripts/evaluation/inverse_design_evaluation.sh

For detailed evaluation methodology and usage, see docs/evaluation/.

Additional Documentation

For comprehensive guides on various features and evaluations:

Usage Guides (docs/guides/):

• How to use formula replacement
• How to resume training
• How to run ablation studies
• How to run hyperparameter sweeps
• Inverse design usage guide
• Text-only inverse design guide
• Text embedding similarity analysis
• Threshold embedding analysis
• UMAP visualization guide
• Zero-shot QA with LLM baselines
• Query materials database

Evaluation Documentation (docs/evaluation/):

• Evaluation results guide

Acknowledgement

• Pytorch implementation of PaiNN

• Self-supervised learning strategies from CLIP-Lite