When to Ensemble: Identifying Token-Level Points for Stable and Fast LLM Ensembling (ICLR 2026)

This is the official code of the paper: When to Ensemble: Identifying Token-Level Points for Stable and Fast LLM Ensembling (ICLR 2026)

Introduction

This repository provides the implementation for SAFE.

Currently, we support UniTE and GaC as ensemble methods. We acknowledge that parts of our code are adapted from the official implementations of UniTE and GaC. We would like to express our sincere gratitude to the authors of these projects!

🛠️ Environment Setup

Ensure the following environment is configured:

• Python: 3.11
• CUDA: 12.4
• PyTorch: 2.5.1
• Transformers: 4.51.3
• FlashAttention: 2.7.4.post1

Install other Python dependencies with:

pip install -r requirements.txt

🚀 Running the Code

This section explains how to run SAFE

---

🧾 Command

1. BBH and MMLU Datasets

For BBH and MMLU datasets, please use run_safe_mmlu.py:

python run_safe_mmlu.py \
    --model_path1 [path_to_model1] \
    --model_path2 [path_to_model2] \
    --model_path3 [path_to_model3] \
    --data_name [data_name] \
    --ensemble [True/False] \
    --ensemble_type [ensemble_type] \
    --output_path [output_path] \
    --sharpen_type [sharpening method] \
    --gamma 5

Other Datasets (GSM8K, MATH500, ARC-Challenge)

For other reasoning tasks, please use run_safe.py:

python run_safe.py \
    --model_path1 [path_to_model1] \
    --model_path2 [path_to_model2] \
    --model_path3 [path_to_model3] \
    --data_name [data_name] \
    --ensemble [True/False] \
    --ensemble_type [ensemble_type] \
    --output_path [output_path] \
    --sharpen_type [sharpening method] \
    --gamma 5

• Drafter Model: Assign to --model_path1.

• Verifier Models: Assign to --model_path2 and --model_path3.

• Two-Model Ensemble: If you are using a two-model ensemble, simply omit --model_path3.

• Path Format: We recommend using official model names. Ensure paths do not end with a trailing slash (/).

Argument	Type	Default	Description
--model_path{1,2,3}	str	-	Path to the model checkpoints for ensemble
--data_name	str	-	Dataset name: gsm8k or math500 or arc or bbh or mmlu-redux
--ensemble	bool	True	Whether to ensemble: True or False
--ensemble_type	str	unite	Type of ensembling: unite or gac
--sharpen_type	str	heur	Method for probability sharpening: heur or geom
--gamma	int	5	Number of draft tokens to be generated
--output_path	str	-	Where outputs will be saved

Citation

If you find this code useful for your research, please consider citing our paper:

@inproceedings{
yun2026when,
title={When to Ensemble: Identifying Token-Level Points for Stable and Fast {LLM} Ensembling},
author={Heecheol Yun and Kwangmin Ki and Jung Hyun Lee and Eunho Yang},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=kfPkF2ACDM}
}