Search papers, labs, and topics across Lattice.
This paper introduces a recursive, sparse mixture-of-experts framework, "Thinking Pixel," integrated into diffusion models to enhance structured reasoning in multimodal text-to-image generation. The approach iteratively refines visual tokens within joint attention layers through sparse selection of specialized neural modules conditioned on visual tokens, timestep, and conditioning information. Experiments on ImageNet, GenEval, and DPG benchmarks demonstrate improved image generation performance compared to standard diffusion models.
Diffusion models can now reason recursively over visual tokens, achieving state-of-the-art image generation performance by dynamically selecting specialized neural modules at each diffusion step.
Diffusion models have achieved success in high-fidelity data synthesis, yet their capacity for more complex, structured reasoning like text following tasks remains constrained. While advances in language models have leveraged strategies such as latent reasoning and recursion to enhance text understanding capabilities, extending these to multimodal text-to-image generation tasks is challenging due to the continuous and non-discrete nature of visual tokens. To tackle this problem, we draw inspiration from modular human cognition and propose a recursive, sparse mixture-of-experts framework integrated into conventional diffusion models. Our approach introduces a recursive component within joint attention layers that iteratively refines visual tokens over multiple latent steps while efficiently sharing parameters via sparse selection of neural modules. At each step, a gating network is devised to dynamically select specialized neural modules, conditioned on the current visual tokens, the diffusion timestep, and the conditioning information. Comprehensive evaluation on class-conditioned ImageNet image generation tasks and additional studies on the GenEval and DPG benchmark demonstrate the superiority of the proposed method in enhancing model image generation performance.