SWIFT

Token-Level Self-Play with Importance-Aware Guidance for Large Language Models

SWIFT (Self-Play Weighted Fine-Tuning) extends self-play alignment with teacher-guided token-level importance weighting.

Overview

SWIFT builds on SPIN and improves token-level learning signals during self-play fine-tuning.
Instead of treating every token equally, SWIFT uses token importance estimated from a stronger teacher model, enabling better alignment and stronger distillation behavior.

Core Idea

• Token-level weighting: focus optimization on more informative tokens.
• Teacher-guided guidance: use a stronger model for token importance instead of direct logits matching.
• Practical tokenizer mapping: transfer token weights across teacher/student tokenizers.

Installation

conda env create -f environment.yml
conda activate SWIFT
pip install -r requirements.txt

Data Preparation

Download:

• SFT data: Link SFT data
• Preference data: Link DPO data

Organize under:

data/
└── Ultrachat200k/
    ├── DPO/
    ├── SFT/
    ├── SPIN/
    └── SWIFT/

Training Pipeline

1) Download Base Checkpoint

python scripts/download.py

2) Run Self-Play Training

bash scripts/_SPIN_full.sh
bash scripts/_SWIFT_full.sh

(_SWIFT_full.sh is the Self-Play Weighted Fine-Tuning pipeline used for SWIFT in this repo.)

Citation

@inproceedings{letoken,
  title={Token-Level Self-Play with Importance-Aware Guidance for Large Language Models},
  author={Le, Tue and Vuong, Hoang Tran and Tran, Quyen and Van, Linh Ngo and Harandi, Mehrtash and Le, Trung},
  booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
  year={2025},
}

License

This project is licensed under the Apache License 2.0. See LICENSE.