KAIST and Yoshua Bengio develop breakthrough to make AI diffusion models 100 times faster

Overview of the inference-time scalable AI diffusion model developed by KAIST researchers Courtesy of KAIST — Overview of the inference-time scalable AI diffusion model developed by a joint research team including KAIST researchers. Courtesy of KAIST.

SEOUL, July 20 (AJP) - A research team at South Korea’s prestigious research institution KAIST, in collaboration with renowned artificial intelligence scholar Yoshua Bengio from the University of Montreal, has developed a technique that significantly accelerates the performance of AI diffusion models while maintaining high output quality. The new method addresses one of the biggest barriers to deploying diffusion models in real-time applications such as robotics and decision-making systems.

KAIST said on Sunday that the project was led by Professor Ahn Sung-jin from KAIST’s School of Computing was carried out through the KAIST-MILA Prefrontal AI Joint Research Center collaboration project. Diffusion models, widely used in AI for tasks like image and text generation, are known for their slow inference speed, which has limited their real-world utility in time-sensitive systems. The research team tackled this issue by introducing an inference approach based on Monte Carlo Tree Search, a strategy often used in game AI to explore decision paths.

By modeling the generation process as a tree and strategically narrowing down the most promising paths, the method was able to produce high-quality outputs using limited computational resources. It also achieved a 100 percent success rate in a large-scale maze-solving task, a benchmark that previous models consistently failed to complete.

The study was presented as a Spotlight paper at the 2025 International Conference on Machine Learning (ICML), held in Vancouver from July. 13 to 19. Spotlight papers are selected from the top 2.6 percent of all submissions, reflecting the technical significance and novelty of the work.

From left Professor Ahn Sung-jin PhD candidate Yoon Jae-sik masters student Cho Hyeon-seo masters student Baek Doo-jin and Professor Yoshua Bengio Courtesy of KAIST — (From left) Professor Ahn Sung-jin, Ph.D. candidate Yoon Jae-sik, master's student Cho Hyeon-seo, master's student Baek Doo-jin, and Professor Yoshua Bengio. Courtesy of KAIST.

While the initial version of the model demonstrated impressive accuracy, its speed remained a challenge. To address this, the team conducted follow-up research to parallelize the tree search algorithm, dramatically boosting its efficiency. The revised model achieved inference speeds up to 100 times faster than previous versions, with no loss in performance.

"This study fundamentally breaks through the high computational cost barrier of conventional diffusion models," said Ahn Sung-jin. "We expect the technique to serve as a core technology across various domains, from intelligent robotics and simulation-based planning to real-time generative AI."

The research was led by Ph.D. candidate Yoon Jae-sik, with contributions from Cho Hyeon-seo, Baek Doo-jin, Bengio, and Ahn. The results were published in two papers. The first, titled Monte Carlo Tree Diffusion for System 2 Planning, was authored by Yoon Jae-sik, Cho Hyeon-seo, Baek Doo-jin, Yoshua Bengio, and Ahn Sung-jin. The second paper, Fast Monte Carlo Tree Diffusion: 100x Speedup via Parallel Sparse Planning, was co-authored by Yoon Jae-sik, Cho Hyeon-seo, Yoshua Bengio, and Ahn Sung-jin. Both were posted on the preprint server arXiv.

The study was supported by the National Research Foundation of South Korea.

Park Sae-jin swatchsjp@ajunews.com