Layout-Corrector: Alleviating Layout Sticking Phenomenon in Discrete Diffusion Model

TL;DR

This work addresses layout-sticking in discrete diffusion models for layout generation by introducing Layout-Corrector, an external Transformer-based correctness assessor that scores per-token harmony and reinitializes low-scoring tokens to the [MASK] state to guide regeneration. By applying the corrector at selected timesteps and masking underperforming tokens, the method steers diffusion models toward more harmonious layouts while enabling control over fidelity-diversity and fast-sampling trade-offs. Extensive experiments on Rico, PubLayNet, and Crello show consistent improvements over state-of-the-art baselines across unconditional and conditional generation tasks, with favorable FID, Alignment, and Max-IoU outcomes and improved qualitative harmony. The approach is model-agnostic within the discrete diffusion family, compatible with LayoutDM, VQDiffusion, and MaskGIT, and offers tunable schedules to balance quality and diversity, albeit with modest additional memory and computation.

Abstract

Layout generation is a task to synthesize a harmonious layout with elements characterized by attributes such as category, position, and size. Human designers experiment with the placement and modification of elements to create aesthetic layouts, however, we observed that current discrete diffusion models (DDMs) struggle to correct inharmonious layouts after they have been generated. In this paper, we first provide novel insights into layout sticking phenomenon in DDMs and then propose a simple yet effective layout-assessment module Layout-Corrector, which works in conjunction with existing DDMs to address the layout sticking problem. We present a learning-based module capable of identifying inharmonious elements within layouts, considering overall layout harmony characterized by complex composition. During the generation process, Layout-Corrector evaluates the correctness of each token in the generated layout, reinitializing those with low scores to the ungenerated state. The DDM then uses the high-scored tokens as clues to regenerate the harmonized tokens. Layout-Corrector, tested on common benchmarks, consistently boosts layout-generation performance when in conjunction with various state-of-the-art DDMs. Furthermore, our extensive analysis demonstrates that the Layout-Corrector (1) successfully identifies erroneous tokens, (2) facilitates control over the fidelity-diversity trade-off, and (3) significantly mitigates the performance drop associated with fast sampling.

Figures (28)

• Figure 1: Intuitive overview of Layout-Corrector. Conventional generative models cannot modify the elements once they have been generated. Layout-Corrector works in conjunction with DDMs to identify inharmonious elements in the generative process and initialize them to enhance regeneration towards a harmonized layout.
• Figure 2: Results of preliminary experiments on Rico test set deka2017rico. (a) While $\bar{\beta}_{t,K}=(K+1)\bar{\beta}_t=\epsilon~(\ll 1)$ is affected by token-sticking, $\bar{\beta}_{t,K} > \epsilon$ alleviates it. (b) The results indicate that LayoutDM can restore the original tokens from [MASK]; however, recovery from regular tokens proves challenging. Please refer to Supp. for further results.
• Figure 3: The details of Layout-Corrector. Top: training procedure of Layout-Corrector, where the pre-trained DDM is fixed. Bottom: sampling process with Layout-Corrector. We execute the generation and correction process in the purple box iteratively.
• Figure 4: Accuracy of detecting erroneous tokens when three tokens are replaced randomly.
• Figure 5: Correctness scores against an extent of layout disruption controlled by various maximum transition steps.