File size: 5,892 Bytes
df21f73
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
832f79f
df21f73
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
43d3eab
df21f73
 
 
 
 
 
 
 
 
 
 
 
43d3eab
df21f73
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
---
license: cc-by-nc-sa-4.0
tags:
  - materials-science
  - property-prediction
  - modular-learning
  - graph-neural-network
  - crystal
datasets:
  - matminer
language:
  - en
library_name: pytorch
pipeline_tag: other
---

# MoMa Hub: Pretrained Modules for Material Property Prediction

<div align="center">

[![arXiv](https://img.shields.io/badge/arXiv-2502.15483-b31b1b.svg)](https://arxiv.org/abs/2502.15483)
[![GitHub](https://img.shields.io/badge/GitHub-MoMa-blue.svg)](https://github.com/Thomaswbt/MoMa)
[![Website](https://img.shields.io/badge/🌐%20Website-MoMa-blue)](https://GenSI-THUAIR.github.io/MoMa/)

</div>

This repository hosts the **18 pretrained full modules** of the MoMa Hub, from the paper:

> **MoMa: A Modular Deep Learning Framework for Material Property Prediction**
>
> Botian Wang, Yawen Ouyang, Yaohui Li, Yiqun Wang, Haorui Cui, Jianbing Zhang, Xiaonan Wang, Wei-Ying Ma, Hao Zhou
>
> *ICLR 2026*

## Model Description

MoMa (**Mo**dular learning for **Ma**terials) is a modular deep learning framework that addresses the diversity and disparity challenges in material property prediction. Instead of forcing all tasks into one shared model, MoMa trains specialized modules on diverse high-resource material tasks and adaptively composes synergistic modules for each downstream scenario.

Each module in this repository is a **full module** β€” a complete GemNet-OC backbone (initialized from the [JMP-L](https://github.com/facebookresearch/JMP) pretrained model) that has been fully fine-tuned on a specific material property prediction task. These modules are designed to be composed via weighted averaging for adaptation to new downstream tasks.

## Modules

This repository contains 18 `.pt` checkpoint files, each trained on a distinct material property prediction task from the [Matminer](https://hackingmaterials.lbl.gov/matminer/) datasets. The modules cover **electronic, thermal, mechanical, and optical** properties across different material databases:

| File | Source | Property | Category |
|------|--------|----------|----------|
| `mp_eform.pt` | Materials Project | Formation Energy | Thermal |
| `mp_bandgap.pt` | Materials Project | Band Gap | Electronic |
| `mp_gvrh.pt` | Materials Project | Shear Modulus (VRH) | Mechanical |
| `mp_kvrh.pt` | Materials Project | Bulk Modulus (VRH) | Mechanical |
| `castelli_eform.pt` | Castelli et al. | Formation Energy | Thermal |
| `jarvis_eform.pt` | JARVIS-DFT | Formation Energy | Thermal |
| `jarvis_bandgap.pt` | JARVIS-DFT | Band Gap (OPT) | Electronic |
| `jarvis_gvrh.pt` | JARVIS-DFT | Shear Modulus (VRH) | Mechanical |
| `jarvis_kvrh.pt` | JARVIS-DFT | Bulk Modulus (VRH) | Mechanical |
| `jarvis_dielectric_opt.pt` | JARVIS-DFT | Dielectric Constant (OPT) | Electronic |
| `n_Seebeck.pt` | Ricci et al. | n-type Seebeck Coefficient | Thermoelectric |
| `n_avg_eff_mass.pt` | Ricci et al. | n-type Average Effective Mass | Thermoelectric |
| `n_e_cond.pt` | Ricci et al. | n-type Electrical Conductivity | Thermoelectric |
| `n_th_cond.pt` | Ricci et al. | n-type Thermal Conductivity | Thermoelectric |
| `p_Seebeck.pt` | Ricci et al. | p-type Seebeck Coefficient | Thermoelectric |
| `p_avg_eff_mass.pt` | Ricci et al. | p-type Average Effective Mass | Thermoelectric |
| `p_e_cond.pt` | Ricci et al. | p-type Electrical Conductivity | Thermoelectric |
| `p_th_cond.pt` | Ricci et al. | p-type Thermal Conductivity | Thermoelectric |

## Architecture

- **Backbone**: GemNet-OC (Large)
- **Initialization**: [JMP-L](https://github.com/facebookresearch/JMP) pretrained checkpoint
- **Module type**: Full module (all backbone parameters fine-tuned)
- **Parameters per module**: ~165M
- **File size per module**: ~615 MB
- **Total repository size**: ~10.8 GB

## Usage

### Download

**Option 1: Using `huggingface_hub` (Python)**

```python
from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="GenSI/MoMa-modules-ICLR",
    repo_type="model",
    local_dir="./hub",
)
```

**Option 2: Using Hugging Face CLI**

```bash
pip install huggingface_hub
hf download GenSI/MoMa-modules-ICLR --repo-type model --local-dir ./hub
```

### Integration with MoMa

After downloading, place the `hub/` directory under the [MoMa codebase](https://github.com/Thomaswbt/MoMa) root:

```
MoMa/
β”œβ”€β”€ hub/
β”‚   β”œβ”€β”€ mp_eform.pt
β”‚   β”œβ”€β”€ mp_bandgap.pt
β”‚   └── ... (18 modules)
β”œβ”€β”€ configs/
β”œβ”€β”€ scripts/
└── ...
```

Then follow the instructions in the [MoMa repository](https://github.com/Thomaswbt/MoMa) to run Adaptive Module Composition and downstream fine-tuning:

```bash
# Adaptive Module Assembly (can be skipped using precomputed results in json/)
bash scripts/extract_embeddings.sh
python scripts/run_knn.py
python scripts/weight_optimize.py

# Downstream Fine-tuning with MoMa
bash scripts/finetune_moma.sh
```

### Loading a Single Module

Each `.pt` file is a standard PyTorch checkpoint containing a `state_dict`:

```python
import torch

ckpt = torch.load("hub/mp_eform.pt", map_location="cpu")
state_dict = ckpt["state_dict"]
```

## Results

MoMa achieves state-of-the-art performance on 17 material property prediction benchmarks, with an average improvement of **14%** over the strongest baseline (JMP fine-tuning). See the full results in our [paper](https://arxiv.org/abs/2502.15483).

## Citation

```bibtex
@article{wang2025moma,
  title={MoMa: A Modular Deep Learning Framework for Material Property Prediction},
  author={Wang, Botian and Ouyang, Yawen and Li, Yaohui and Wang, Yiqun and Cui, Haorui and Zhang, Jianbing and Wang, Xiaonan and Ma, Wei-Ying and Zhou, Hao},
  journal={arXiv preprint arXiv:2502.15483},
  year={2025}
}
```

## License

This work is licensed under a [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License](http://creativecommons.org/licenses/by-nc-sa/4.0/).