CrysTune: Crystal Generation via Fine-Tuning of Large Language Models on Wyckoff Representations

Motivation

• Known: 기존 연구는 Crystal-Text-LLM처럼 문자열 기반 표현으로 LLM을 fine-tuning하거나, CrystalFormer처럼 Wyckoff position을 인코딩한 custom transformer를 학습시켜 결정의 대칭성을 반영하려 시도해왔다.
• Gap: Crystal-Text-LLM은 space-group이나 site-symmetry 같은 구조적 inductive bias를 명시적으로 학습에 반영하지 않고, CrystalFormer는 대칭성을 반영하지만 pretrained LLM의 지식을 활용하지 못하며, 기존 LLM 기반 접근은 next-token prediction만으로 학습되어 3D 결정 구조에 내재된 대칭성을 충분히 포착하지 못한다.
• Why: 대칭성과 주기성을 고려한 inductive bias를 효과적으로 결정 생성 모델에 통합하는 것은 제한된 재료 데이터 환경에서 신규 안정 결정 구조를 효율적으로 탐색하는 데 핵심적이며, LLM의 사전지식과 파라미터 효율적 fine-tuning을 결합하면 특수 아키텍처 설계 없이도 이를 달성할 잠재력이 있다.
• Approach: CrysTune은 LLaMA-2 chat 모델을 Wyckoff string representation과 canonicalization, template prediction이라는 두 보조 과제로 LoRA fine-tuning한 후, 이를 GRPO 기반 RL의 초기 정책으로 사용해 안정성·유효성·고유성·신규성·다양성을 최적화하는 2단계 파이프라인을 제안한다.

Achievement

1. Wyckoff 기반 LLM fine-tuning: CrystalFormer에서 채택한 Wyckoff-based crystal representation을 LLaMA-2 chat 모델의 tokenizer에 직접 적용해 fine-tuning함으로써 custom architecture 없이 대칭성 정보를 학습에 반영했다.
2. 보조 과제 도입을 통한 대칭 학습 강화: canonicalization과 template prediction이라는 두 신규 보조 과제를 도입해 생성 과제만 학습한 모델이나 canonicalization만 추가한 모델보다 우수한 성능과 안정성 관련 지표 개선을 달성했다.
3. RL fine-tuning을 통한 성능 개선: Park & Walsh (2025)의 보상 함수를 채택해 GRPO 기반 RL fine-tuning을 수행한 결과, 적절한 학습 스텝 수에서 유효성·안정성·다양성이 크게 향상되었으나 novelty는 크게 감소하는 trade-off를 확인했다.
4. 속성 조건부 생성 검증: crystal system 조건화 실험을 통해 RL로 학습된 정책이 목표 조건에 부합하는 결정을 더 높은 비율로 생성함을 입증했다.

How

• 무기 결정 데이터셋을 space group number, Wyckoff symbol, 원소, 분수 좌표, lattice parameter로 구성된 Wyckoff string representation으로 변환한다.
• LLaMA-2 7B-chat 모델을 LoRA로 fine-tuning하며, task-specific prompt를 사용해 세 가지 과제(Task 1: Generation, Task 2: Canonicalization, Task 3: Template Prediction)를 함께 학습한다.
• SFT로 학습된 CrysTune 모델을 초기 정책으로 삼아 Group Relative Policy Optimization(GRPO) 기반 RL fine-tuning을 수행하며, group rollout을 통해 validity, diversity, stability, property를 반영한 reward를 계산한다.
• 생성된 구조에 대해 strict validity checks, multiplicity correction, fixed coordinate correction, orbit expansion 등의 post-processing을 적용해 유효한 crystal structure를 얻는다.
• LeMat-GenBench를 사용해 validity, novelty, diversity, stability 등 핵심 지표로 baseline 및 ablation 모델을 평가한다.

Originality

• Wyckoff representation을 pretrained LLM의 tokenizer에 직접 적용해 custom architecture 없이 symmetry-aware crystal generation을 수행한 최초 시도 중 하나이다.
• canonicalization과 template prediction이라는 두 가지 novel auxiliary task를 도입해 LLM이 결정 대칭성을 명시적으로 학습하도록 유도한 점이 독창적이다.
• SFT로 학습된 LLM 기반 crystal generator를 GRPO를 통해 RL fine-tuning하여 다중 목표(안정성, 유효성, 고유성, 신규성, 다양성)를 동시에 최적화하는 파이프라인을 제시했다.

Limitation & Further Study

• RL fine-tuning 과정에서 novelty가 크게 감소하고 mode collapse가 발생하는 문제가 있어, 안정성·유효성과 novelty·diversity 간의 근본적인 trade-off 해결이 향후 과제로 남아있다.
• 논문 발췌 내용에는 정량적 비교표나 세부 실험 설정이 충분히 드러나지 않아 baseline 대비 개선폭의 통계적 유의성 및 재현성 검증이 필요하다.
• 7B 규모의 단일 LLM 계열(LLaMA-2)에 대해서만 실험이 이루어져, 다른 아키텍처나 더 큰 모델로의 일반화 가능성에 대한 추가 검증이 요구된다.
• 후속 연구로 novelty-diversity와 stability-validity 간 trade-off를 완화할 수 있는 보상 설계나 다목적 최적화 기법 탐구가 필요하다.

Evaluation

총평: Wyckoff representation과 보조 과제를 결합한 LLM fine-tuning 및 GRPO 기반 RL을 통해 symmetry-aware crystal generation의 새로운 방향을 제시한 흥미로운 연구이나, novelty-diversity trade-off와 정량적 검증의 충실도가 향후 개선 과제로 남아있다.