Ant Group's UBTree improves speculative decoding: parallel screening of draft trees enables large models to verify more tokens at once.

Beating AI News Flash: Ant Bailing has open-sourced UBTree, a speculative decoding solution for large models.

Speculative decoding is akin to having a small model "draft" ahead of time, which is then verified by a large model in batches of multiple tokens, reducing the computational rounds needed to generate tokens one by one.

However, drafts also suffer from hit rate issues. The earlier the small model guesses wrong, the less content can be directly adopted afterward. Previous technologies prepared multiple candidate paths simultaneously, forming a draft tree, and pruned branches unlikely to be accepted.

UBTree improves the filtering method. It first predicts in parallel the tokens that may appear at subsequent positions, then uses a lightweight model to judge whether adjacent tokens can connect reasonably. Different combinations can be scored simultaneously without repeatedly running the model along each path. Promising branches are then selected to form a draft tree for verification by the large model.

The team also adjusted the training method of the filtering model, enabling it to not only focus on the highest-probability answer but also identify alternative paths worth retaining, giving the large model the opportunity to accept more tokens per round.

In seven tests on Qwen3-4B, the average speedup ratio using the DFlash 2 selector to build trees was 5.95x. After switching to the UBTree selector and adjusting the training method, it increased to 6.88x, equivalent to a further improvement of approximately 16% over the original scheme.