crystal-llm-representations

LoRA adapters and training logs accompanying the paper:

Input Representation and Token-Level Confidence for LLMs in Materials Property Prediction

Paper under review.

Code: github.com/shu-ozawa/crystal-llm-representations

The adapters of the submitted version (rank 32, 1B and 8B models, formation energy and bandgap) remain available under the tag v1-submission.

Contents

135 LoRA adapters, one for every combination of 3 properties, 5 input representations, 3 model sizes, and 3 random seeds, with the training log of each.

lora_adapters/{property}/seed{seed}/{size}-{input}-r{rank}/
    adapter_config.json
    adapter_model.safetensors
training_logs/{property}/seed{seed}/{size}-{input}-r{rank}.log

Field	Values
property	formation_energy_per_atom, band_gap (Materials Project), bulk_modulus_kv (JARVIS-DFT), from LLM4Mat-Bench
seed	43, 44, 45
size, rank	1B with r256 (Llama-3.2-1B-Instruct), 3B with r256 (Llama-3.2-3B-Instruct), 8B with r128 (Llama-3.1-8B-Instruct)
input	formula (Composition), descr1st (Crystal Summary), descr_last (Local Environment), description (Full Description), cif (CIF)

For example, lora_adapters/band_gap/seed43/8B-description-r128 is the Llama-3.1-8B-Instruct adapter for bandgap with the Full Description input and seed 43.

Training

LoRA is applied to the query and value projections with a scaling factor of α = 2r and a dropout of 0.05. Each adapter is trained for 6 epochs with a learning rate of 1e-4 (cosine schedule, warmup ratio 0.08), AdamW with a weight decay of 0.01, an effective batch size of 32, and a maximum sequence length of 8,192 tokens, with the loss computed over the response only. The checkpoint with the lowest validation loss is kept. The chat template date is fixed to 26 Jul 2024.

Usage

The prompts are built by src/prepare_dataset_lora.py in the code repository, and src/run_inference_with_hidden.py runs an adapter on a prepared dataset. The models answer in the form

{formation_energy_per_atom: -2.746 eV/atom}
{band_gap: 0.0 eV}
{bulk_modulus_kv: 33.13}

To download a single adapter:

from huggingface_hub import snapshot_download

name = "lora_adapters/band_gap/seed43/8B-description-r128"
root = snapshot_download("ozashu/crystal-llm-representations", allow_patterns=f"{name}/*")
adapter_dir = f"{root}/{name}"

python src/run_inference_with_hidden.py \
    --dataset_dir data/prepared/mp/band_gap/description \
    --adapter_dir /path/to/adapter_dir \
    --output_dir outputs/mp/eval_description/band_gap/Bg8B-description-r128-seed43/test \
    --max_new_tokens 512 \
    --max_input_length 8192 \
    --date_string "26 Jul 2024"

Training logs

Each log is the console output of one training run, including the validation loss after every epoch. Paths and host names in the logs refer to the machines the runs were carried out on.

• Formation energy, seed 43, 3B with Crystal Summary, Full Description, and CIF: these runs were first trained in the rank sweep, whose r256 runs are the adopted adapters, so their logs refer to that sweep.
• Bandgap, seed 43, 8B with Full Description and Local Environment: these runs were interrupted and resumed from checkpoints several times, and the console output was saved only for some segments. Each log joins the saved segments, marked by ===== lines, in the order they were saved, so the step count goes back at each resume. The end of training is missing: the Full Description log stops at step 18,718 and the Local Environment log at step 11,773 of 23,460. The adapters are complete.