From Global to Factor-Wise Expert Composition in Discrete Diffusion Models

Haozhe Huang, Yudong Xu, Abhijoy Mandal, Alan Aspuru-Guzik
Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, PMLR 337:2223-2243, 2026.

Abstract

Discrete diffusion models offer a powerful framework for solving complex reasoning tasks, particularly through compositional generation, which combines multiple pre-trained experts to generalize beyond their individual training data. Recent theoretical corrections introduce time-dependent mixing weights to better align composed diffusion dynamics with the intended target. However, these methods are fundamentally limited by working on a per-sample basis, treating each generated state monolithically and ignoring the potential spatial or functional specializations of different experts. In this work, we address this limitation by proposing FactorDiff - a factor-wise composition framework for diffusion models. We posit that samples can be further decomposed into smaller factors, and propose a sampling process that dynamically routes each factor to the most relevant expert. We instantiate this framework with spatial/pixel-level compositions and validate it on the {ARC-AGI} benchmark, demonstrating that simple factor-specific routing consistently outperforms complex global scalar weighting schemes on tasks that require logical consistency and spatial disentanglement.

Cite this Paper

BibTeX
@InProceedings{pmlr-v337-huang26a, title = {From Global to Factor-Wise Expert Composition in Discrete Diffusion Models}, author = {Huang, Haozhe and Xu, Yudong and Mandal, Abhijoy and Aspuru-Guzik, Alan}, booktitle = {Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence}, pages = {2223--2243}, year = {2026}, editor = {Perković, Emilija and Malinsky, Daniel}, volume = {337}, series = {Proceedings of Machine Learning Research}, month = {17--21 Aug}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v337/main/assets/huang26a/huang26a.pdf}, url = {https://proceedings.mlr.press/v337/huang26a.html}, abstract = {Discrete diffusion models offer a powerful framework for solving complex reasoning tasks, particularly through compositional generation, which combines multiple pre-trained experts to generalize beyond their individual training data. Recent theoretical corrections introduce time-dependent mixing weights to better align composed diffusion dynamics with the intended target. However, these methods are fundamentally limited by working on a per-sample basis, treating each generated state monolithically and ignoring the potential spatial or functional specializations of different experts. In this work, we address this limitation by proposing FactorDiff - a factor-wise composition framework for diffusion models. We posit that samples can be further decomposed into smaller factors, and propose a sampling process that dynamically routes each factor to the most relevant expert. We instantiate this framework with spatial/pixel-level compositions and validate it on the {ARC-AGI} benchmark, demonstrating that simple factor-specific routing consistently outperforms complex global scalar weighting schemes on tasks that require logical consistency and spatial disentanglement.} }
Endnote
%0 Conference Paper %T From Global to Factor-Wise Expert Composition in Discrete Diffusion Models %A Haozhe Huang %A Yudong Xu %A Abhijoy Mandal %A Alan Aspuru-Guzik %B Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence %C Proceedings of Machine Learning Research %D 2026 %E Emilija Perković %E Daniel Malinsky %F pmlr-v337-huang26a %I PMLR %P 2223--2243 %U https://proceedings.mlr.press/v337/huang26a.html %V 337 %X Discrete diffusion models offer a powerful framework for solving complex reasoning tasks, particularly through compositional generation, which combines multiple pre-trained experts to generalize beyond their individual training data. Recent theoretical corrections introduce time-dependent mixing weights to better align composed diffusion dynamics with the intended target. However, these methods are fundamentally limited by working on a per-sample basis, treating each generated state monolithically and ignoring the potential spatial or functional specializations of different experts. In this work, we address this limitation by proposing FactorDiff - a factor-wise composition framework for diffusion models. We posit that samples can be further decomposed into smaller factors, and propose a sampling process that dynamically routes each factor to the most relevant expert. We instantiate this framework with spatial/pixel-level compositions and validate it on the {ARC-AGI} benchmark, demonstrating that simple factor-specific routing consistently outperforms complex global scalar weighting schemes on tasks that require logical consistency and spatial disentanglement.
APA
Huang, H., Xu, Y., Mandal, A. & Aspuru-Guzik, A.. (2026). From Global to Factor-Wise Expert Composition in Discrete Diffusion Models. Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, in Proceedings of Machine Learning Research 337:2223-2243 Available from https://proceedings.mlr.press/v337/huang26a.html.

Related Material