---
license: mit
language:
- en
tags:
- molecules
- chemistry
- molecular-understanding
- graph-llm
- llama
- instruction-tuning
pipeline_tag: text-generation
base_model: unsloth/Llama-3.1-8B-Instruct
---

# EDT-Former: Full Model (Stage 2)

The full **EDT-Former** model (encoder + Llama-3.1-8B-Instruct), as described in the ICLR 2026 paper:

> **Entropy-Guided Dynamic Tokens for Graph-LLM Alignment in Molecular Understanding**
> Zihao Jing, Qiuhao Zeng, Ruiyi Fang, Yan Sun, Boyu Wang, Pingzhao Hu
> *ICLR 2026* · [Paper](https://www.arxiv.org/abs/2602.02742) · [Code](https://github.com/selmiss/DQ-Former)

## Model Description

EDT-Former aligns molecular graphs with a frozen LLM backbone (Llama-3.1-8B-Instruct) via an entropy-guided dynamic token connector. Key properties:

- **No LLM backbone fine-tuning** (only the embedding layer and connector are trained) — computationally efficient
- **Entropy-guided dynamic token selection** preserves both local (substructural) and global molecular features
- **State-of-the-art** on MoleculeQA, Mol-Instructions (forward reaction, retrosynthesis, reagent prediction, mol design, open QA), TDC, and MoleculeNet benchmarks

This Stage 2 checkpoint (~16 GB) is the final instruction-tuned model ready for downstream molecular QA tasks.

## Usage

```python
# 1. Clone the repo and set up the environment
git clone https://github.com/selmiss/DQ-Former.git
cd DQ-Former
conda env create -f environment.yml
conda activate edtformer

# 2. Configure paths in local.env.sh
cp env.sh local.env.sh
# Edit local.env.sh: set BASE_DIR, DATA_DIR, CHECKPOINT_DIR
source local.env.sh

# 3. Download the model
from huggingface_hub import snapshot_download
snapshot_download("zihaojing/EDT-Former-model", local_dir="checkpoints/edt_former_s2_large/final_model")

# 4. Run inference (example: forward reaction prediction)
bash scripts/qa/mol_forward.sh
```

## Downstream Task Scripts

All evaluation scripts are in `scripts/qa/`. Example tasks:

| Task | Script |
|------|--------|
| Forward Reaction Prediction | `scripts/qa/mol_forward.sh` |
| Retrosynthesis | `scripts/qa/retrosynthesis.sh` |
| Reagent Prediction | `scripts/qa/reagent_prediction.sh` |
| Molecule Design | `scripts/qa/mol_design.sh` |
| Open-ended QA | `scripts/qa/open_question.sh` |

## Training Details

| Setting | Value |
|---------|-------|
| LLM backbone | Llama-3.1-8B-Instruct (frozen) |
| Stage 1 encoder | [zihaojing/EDT-Former-encoder](https://huggingface.co/zihaojing/EDT-Former-encoder) |
| Training data | [zihaojing/EDT-Former-sft-data](https://huggingface.co/datasets/zihaojing/EDT-Former-sft-data) |
| Epochs | 2 |
| Learning rate | 1e-4 (cosine) |
| Batch size | 4 × 8 grad accum = effective 32 |
| Precision | BF16 |

## Related Resources

| Resource | Link |
|----------|------|
| Pretrain Data | [zihaojing/EDT-Former-pretrain-data](https://huggingface.co/datasets/zihaojing/EDT-Former-pretrain-data) |
| SFT Data | [zihaojing/EDT-Former-sft-data](https://huggingface.co/datasets/zihaojing/EDT-Former-sft-data) |
| Encoder (Stage 1) | [zihaojing/EDT-Former-encoder](https://huggingface.co/zihaojing/EDT-Former-encoder) |
| Code | [selmiss/DQ-Former](https://github.com/selmiss/DQ-Former) |

## Citation

```bibtex
@inproceedings{jing2026edtformer,
  title={Entropy-Guided Dynamic Tokens for Graph-LLM Alignment in Molecular Understanding},
  author={Jing, Zihao and Zeng, Qiuhao and Fang, Ruiyi and Sun, Yan and Wang, Boyu and Hu, Pingzhao},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2026}
}
```