Milyte

A Multiscale, Physics-Native Deep Learning Framework
from Molecule to Cell for Electrolyte Property Prediction and Formulation Design

---

Milyte connects four physical scales of the battery electrolyte — molecular solvation structure, mesoscopic aggregation, bulk transport properties, and device-level performance — within a single differentiable model. A Transformer backbone is pretrained on molecular-dynamics-derived radial distribution functions (RDFs) under a triple objective (RDF reconstruction, stepwise desolvation free-energy prediction, SE(3)-equivariant denoising), then finetuned into four downstream property predictors and inverted for target-conditioned formulation design (GALO).

Data and model policy. All pretrained weights and datasets are hosted separately on Hugging Face:

🤗 huggingface.co/xiaomi-research/milyte

Highlights

🧬 Physics-native pretraining	RDF reconstruction + stepwise ΔG_des prediction + SE(3)-equivariant denoising, learning the local Li⁺ solvation shell
🔗 One backbone, four properties	Ionic conductivity (κ), viscosity (η), desolvation free energy (ΔG), dual-temperature DCR
⚗️ Up to 7-component formulations	Mixture-invariant composition encoding + Arrhenius-inspired temperature kernel
🎯 Inverse design (GALO)	Gradient ascent over the composition simplex + local refinement — target property → candidate recipe
✅ Reproducible evaluation	Consistent 5-fold cross-validation across all property models

Repository layout

milyte/
├── README.md                  # This document
├── LICENSE                    # Apache License 2.0
├── environment.yml            # Conda environment definition (recommended)
├── requirements.txt           # pip dependencies (shared by both stages)
├── data/                      # Local mount point, populated from Hugging Face (not tracked in Git)
│   ├── checkpoints/           # Pretrained RDF backbone: best_model.pt, config.json, vocab/mol files
│   ├── pretrain/              # Pretraining data: pretrain_rdf_{train,val,test}.csv
│   └── downstream/            # Finetuning/evaluation data and split artifacts
│       ├── conductivity.csv / viscosity.csv / free_energy.csv / dcr.csv
│       └── mol_columns.json / ion_pair_vocab.json
├── src/
│   ├── pretrain/              # Stage I: RDF pretraining
│   │   └── train_config.yaml  # Pretraining config example
│   └── downstream/            # Stage II: downstream finetuning + evaluation
│       └── models/
│           ├── common/                           # Shared base: components / data / losses / arch / metrics
│           ├── rdf_backbone.py                   # RDF pretrained backbone + pretraining entry
│           ├── conductivity_predictor.py         # Conductivity prediction
│           ├── free_energy_predictor.py          # Desolvation free energy prediction
│           ├── viscosity_predictor.py            # Viscosity prediction
│           ├── dcr_dual_temperature_predictor.py # Dual-temperature DCR prediction
│           ├── train_dcr_dual_temperature.py     # DCR training entry
│           ├── recipe_inference.py               # Recipe → property inference helper
│           └── inverse_design.py                 # Inverse design (GALO): recipe for a target property
└── tests/
    └── test_smoke.py          # Smoke tests: existence and path resolution of weights/data/config

Data and model availability

The data/ directory shown above is the expected local mount point, populated by the user after downloading from Hugging Face.

🤗 huggingface.co/xiaomi-research/milyte

This includes the pretrained RDF backbone, pretraining RDF sequences (>1.8M ion-pair RDFs across 126 MD-simulated electrolyte systems), downstream experimental datasets, and 5-fold finetuned checkpoints for each property.

pip install huggingface_hub
huggingface-cli download xiaomi-research/milyte --local-dir ./data

Adjust target subpaths as needed so the downloaded files land in data/checkpoints, data/pretrain, and data/downstream as expected by the code.

Environment setup

# Recommended: Conda environment
conda env create -f environment.yml
conda activate electrolyte

For a pip-only setup, run pip install -r requirements.txt.

GPU-enabled PyTorch (optional, after activating the environment):

pip install --upgrade torch torchvision --index-url https://download.pytorch.org/whl/cu124

Stage I — physics-native RDF pretraining

RDF pretraining data originates from GROMACS 2024 molecular dynamics simulations (gmx rdf) across 126 electrolyte systems. The pretraining objective jointly combines RDF reconstruction, stepwise desolvation free-energy prediction, and SE(3)-equivariant coordinate denoising. Processed sequences are distributed via Hugging Face as:

data/pretrain/pretrain_rdf_train.csv   # training set (RDF sequences)
data/pretrain/pretrain_rdf_val.csv     # validation set
data/pretrain/pretrain_rdf_test.csv    # test set

To pretrain the backbone from scratch:

python src/downstream/models/rdf_backbone.py

Or download the pretrained best_model.pt from Hugging Face and use it directly as the downstream backbone, skipping Stage I entirely.

Stage II — downstream property prediction

All property models share the datasets, losses, component features, and architecture classes in src/downstream/models/common/, and load the pretrained backbone via rdf_backbone.RDFPretrainModel.

Conductivity / free energy / viscosity

cd src/downstream
python models/conductivity_predictor.py    # Train a 5-fold conductivity model
python models/viscosity_predictor.py       # Train a 5-fold viscosity model
python models/free_energy_predictor.py     # Train a 5-fold desolvation free energy model

Dual-temperature DCR

python models/train_dcr_dual_temperature.py --csv ../data/downstream/dcr.csv

Inverse design — GALO (target conductivity → recipe)

src/downstream/models/inverse_design.py inverts the conductivity model with a gradient-ascent-with-local-optimization (GALO) strategy: given a target conductivity, it optimizes the continuous component fractions via Adam over the composition simplex, then optionally refines the optimum with a local grid search.

cd src/downstream/models
python inverse_design.py \
  --checkpoint outputs/conductivity_enhanced_fold1.pth \
  --pretrained ../../../data/checkpoints/best_model.pt \
  --mol-columns ../../../data/downstream/mol_columns.json \
  --target 15 \
  --local-refine

Key options (see python inverse_design.py -h for the full list):

Flag	Default	Description
--target	(required)	Target conductivity, same units as the labels (mS/cm)
--num-restarts	8	Parallel random initial points; keeps the one closest to target
--steps / --lr	500 / 0.05	Adam optimization schedule over the softmax logits
--mol-count-scale	1.0	mol_counts = scale * mol_fractions; set to match the training mol_counts scale (see note below)
--prior-weight / --init-fractions	0.0 / None	Optional MSE prior pulling the solution toward a known recipe
--local-refine	off	Local grid over pairs (i, j) via f_i ± Δf, f_j ∓ Δf, picking the recipe closest to target
--output-order-csv	data/downstream/conductivity.csv	Key order of mol_fractions / recipe follows this file's header columns
--benchmark-speedup	off	Compare inverse-design wall clock vs. an exhaustive discretized grid screen

Note on --mol-count-scale: training uses raw CSV amounts summing to ~100; the default --mol-count-scale 1.0 keeps mol_counts = mol_fractions (sum = 1), which is out-of-distribution. Set it to ~100 (or your training total) to match the forward model.

Testing

pytest tests/   # Checks existence and path resolution of the two-stage weights/data/config

The smoke tests assume the Hugging Face assets have already been downloaded into data/; without them, data/weight presence checks will fail by design.

Citation

If you use Milyte in your research, please cite our paper (citation details to be added upon publication).

License

Code in this repository is released under the Apache License 2.0. Model weights and datasets are hosted separately on Hugging Face and may be subject to distinct terms — see the model card for details.