논문 | OpenReview | 포스터
이 저장소는 "EquiformerV2: 고차원 표현으로 확장하기 위한 개선된 등변환 트랜스포머" (ICLR 2024) 작업의 공식 PyTorch 구현을 담고 있습니다.
저희는 OC20 S2EF-2M 및 S2EF-All+MD 데이터셋에서 기본 모델 설정을 학습시키는 코드를 제공합니다.
또한, EquiformerV2는 OCP repository에 통합되었으며 Open Catalyst demo에서 사용되었습니다.
저희의 후속 작업에서는 BERT와 유사한 자기 지도 학습(self-supervised learning)을 3D 원자 시스템에 일반화할 수 있음을 발견했으며, 이를 DeNS (Denoising Non-Equilibrium Structures)라고 명명했습니다. 이는 에너지 및 힘 예측에서 EquiformerV2의 성능을 향상시킵니다.
더 자세한 내용은 논문과 코드를 참조해 주십시오.
차세대 모델인 EquiformerV3는 여기에서 확인할 수 있습니다.
저희는 효율성, 표현력(expressivity), 일반성을 더욱 개선하여 직접 및 기울기 예측뿐만 아니라 고차 미분(higher-order derivatives)이 필요한 작업에서도 최첨단(state-of-the-art) 결과를 달성했습니다.
<p align="center">
<img src="https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/fig/equiformer_v2.png" alt="photo not available" width="98%" height="98%">
</p>
<p align="center">
<img src="https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/fig/equiformer_v2_speed_accuracy_tradeoffs.png" alt="photo not available" width="98%" height="98%">
</p>
⚠️ [IMG:N] 형식 토큰은 이미지 placeholder 입니다. 번역하지 말고 원래 위치에 그대로 유지하세요.
<p align="center">
<img src="https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/fig/equiformer_v2_oc20_results.png" alt="photo not available" width="98%" height="98%">
</p>
<p align="center">
<img src="https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/fig/equiformer_v2_adsorbml_results.png" alt="photo not available" width="98%" height="98%">
</p>
<p align="center">
<img src="https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/fig/equiformer_v2_oc22_results.png" alt="photo not available" width="98%" height="98%">
</p>
목차
- 환경 설정
- 변경 로그
- 학습
- 파일 구조
- 체크포인트
- 인용
- 감사
환경 설정
환경 설정 (Environment)
환경 설정을 위해서는 여기를 참조하세요.
OC20
OC20 S2EF 데이터셋은 해당 GitHub 저장소의 지침을 따라 다운로드할 수 있습니다.
예를 들어, 다음 명령어를 실행하여 OC20 S2EF-2M 데이터셋을 다운로드할 수 있습니다:
cd ocp
python scripts/download_data.py --task s2ef --split "2M" --num-workers 8 --ref-energy
또한 학습을 실행하려면 `
-
OC20 S2EF-2M 데이터셋에서 다음 명령어를 실행하여 EquiformerV2를 학습시킵니다:
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@[email protected]
위 스크립트는 각 노드에 8개의 GPU를 사용하여 총 2개 노드를 사용합니다.
만약 가져오기 오류(import error)가 발생한다면, ocp/ocpmodels/common/utils.py가 수정되지 않았을 수 있습니다. 자세한 내용은 여기를 따르십시오.
또한 1개 노드의 8개 GPU에서 학습을 실행할 수도 있습니다:
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@[email protected]
-
OC20 S2EF-All+MD에서 다음 명령어를 실행하여 **EquiformerV2 (153M)**를 학습시킵니다:
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@20_L@6_M@3_splits@[email protected]
위 스크립트는 각 노드에 8개의 GPU를 사용하여 총 16개 노드를 사용합니다.
-
OC20 S2EF-All+MD에서 다음 명령어를 실행하여 **EquiformerV2 (31M)**를 학습시킵니다:
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@8_L@4_M@2_splits@[email protected]
위 스크립트는 각 노드에 8개의 GPU를 사용하여 총 8개 노드를 사용합니다.
-
에너지 및 힘 예측 성능을 더욱 향상시키기 위해 보조 과제(auxiliary task)로 DeNS (Denoising Non-Equilibrium Structures)와 함께 EquiformerV2를 학습시킬 수 있습니다. 자세한 내용은 코드를 참조하십시오.
파일 구조
nets에는 OC20을 위한 다양한 네트워크 아키텍처 코드가 포함되어 있습니다.
scripts에는 OC20에서 모델 학습에 필요한 스크립트가 포함되어 있습니다.
main_oc20.py는 학습, 평가 및 이완(relaxation) 실행을 위한 코드입니다.
oc20/trainer에는 힘 트레이너(force trainer) 코드와 일부 유틸리티 함수가 포함되어 있습니다.
oc20/configs에는 S2EF를 위한 설정 파일이 포함되어 있습니다.
체크포인트
S2EF-2M 데이터셋으로 30 에포크 동안 학습된 EquiformerV2 체크포인트, S2EF-All+MD에 대해 학습된 EquiformerV2 (31M), 그리고 S2EF-All+MD에 대해 학습된 EquiformerV2 (153M)을 제공합니다.
이 저장소가 도움이 되었다면 아래 연구들을 인용해 주십시오:
- EquiformerV2:
@inproceedings{
equiformer_v2,
title={{EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations}},
...
```
- eSCN:
@inproceedings{
escn,
title={{Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs}},
...
```
- Equiformer:
@inproceedings{
equiformer,
title={{Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs}},
...
```
- OC20 데이터셋:
@article{
oc20,
author = {Chanussot*, Lowik and Das*, Abhishek and Goyal*, Siddharth and Lavril*, Thibaut and Shuaibi*, Muhammed and Riviere, Morgane and Tran, Kevin and Heras-Domingo, Javier and Ho, Caleb and Hu, Weihua and Palizhati, Aini and Sriram, Anuroop and Wood, Brandon and Yoon, Junwoong and Parikh, Devi and Zitnick, C. Lawrence and Ulissi, Zachary},
...
```
질문은 Yi-Lun Liao ([email protected])에게 문의하십시오.
감사의 말씀
저희의 구현은 PyTorch, PyG, e3nn, timm, ocp, Equiformer를 기반으로 합니다.