zihaojing/EDT-Former-encoder

EDT-Former Encoder (Stage 1)

The pretrained EDT-Former encoder from 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 · Code

Model Description

The EDT-Former encoder is a Dual Q-Former that bridges molecular graphs and language. It uses:

• Entropy-guided dynamic token selection to focus on informative molecular patches
• BRICS fragment IDs for substructural awareness
• Cross-attention over graph node features to generate a token sequence aligned with text

This Stage 1 checkpoint (~699 MB) is trained on the PubChem pretraining corpus and is used to initialize Stage 2 (full model) training.

Architecture config:

• num_query_tokens: 32
• embed_dim: 512
• cross_attention_freq: 1
• num_layers: 8 (blending module)
• num_heads: 8

Usage

Use this checkpoint as the Stage 1 initialization for Stage 2 fine-tuning:

# configs/stage2_dqw2d/model_config.yaml
stage1_path: path/to/EDT-Former-encoder/model.safetensors

Or download and use directly:

from huggingface_hub import snapshot_download

snapshot_download("zihaojing/EDT-Former-encoder", local_dir="checkpoints/edt_former_s1_large/final_model")

To reproduce Stage 1 training from scratch:

# Set up environment first (see repo README)
bash scripts/training/pretraining.sh

Related Resources

Resource	Link
Pretrain Data	zihaojing/EDT-Former-pretrain-data
SFT Data	zihaojing/EDT-Former-sft-data
Full Model (Stage 2)	zihaojing/EDT-Former-model
Code	selmiss/DQ-Former

Citation

@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}
}