LayoutDM: Discrete Diffusion Model for Controllable Layout Generation

In plain words

A layout is the arrangement of boxes on a page or screen: each box has a kind (text, image, button), a position and a size. LayoutDM learns to generate such arrangements by starting from a messy version and cleaning it up step by step. Each box is written as a short row of tokens, one for its kind and one for each coordinate, rounded to a fixed set of values. A diffusion model is trained to undo random damage to those tokens, where the damage is hiding or swapping them.

The useful part is what you can ask for afterwards. You can fix some boxes or some of their properties and let the model fill in the rest, and you can nudge the result toward softer goals, such as keeping a box near where a rough sketch put it, or keeping two boxes in a given relation. None of these needs a separate model, because the constraints are applied while the model is running.

The point is that one model can cover several layout tasks that used to need one model each.

TL;DR

• LayoutDM is a discrete diffusion model over layouts quantized to tokens. Each element is a category plus four quantized geometry tokens; layouts are padded to 25 elements §3.2 §4.1.
• Corruption uses a separate transition matrix for each attribute, and a 4-layer Transformer encoder predicts clean tokens §3.2 §4.4.
• Conditions are applied only at inference: hard masking for known fields and logit adjustment for weak priors such as refinement and relations §3.3.
• Author-reported: LayoutDM outperforms task-agnostic baselines in most cases §1. Unconditional Rico FID is 6.65 tab. 2.
• The official code is Apache-2.0 and deviates from the paper in several places (see Paper vs code).

Problem

A single layout model should handle unconditional generation, conditioning on categories or sizes, completing a partial layout, refining noisy boxes, and satisfying relations between elements §1 §4.3. Autoregressive layout decoders fix an order of elements and attributes, so they fail when a condition disagrees with that order §1. Non-autoregressive models can condition on fields, but they cannot complete a layout because they need the number of elements in advance §2. Standard discrete diffusion treats all tokens as interchangeable, so it can corrupt a category into a coordinate token; the paper's layout tokens form disjoint groups per attribute §3.2.

Method

Representation

• An element is a category c and a box \bm{b} \in [0,1]^4 with center, width and height. Each value is quantized to one of B bins; the layout becomes a sequence of (c, x, y, w, h) tokens §3.2.
• Elements are shuffled at random and flattened, to avoid depending on a fixed order §3.2. The official code applies this as a RandomOrder transform hydra_configs.py.
• A [PAD] token fills unused slots. With at most M = 25 elements, each layout has 5M tokens §3.2 §4.1. Layouts with more elements are discarded rico.py publaynet.py.
• Geometry bins are quantized with KMeans fitted per attribute on the training data §3.2. The paper does not give B. The official code defaults to 32 bins hydra_configs.py and fits the clusters with scikit-learn's KMeans clustering_coordinates.py.
• Vocabulary: categories, then bins, then [PAD] and [MASK]. The official LayoutDM config uses separate bins for x, y, w and h, so the total is C + 4B + 2 app. A.1 layout_tokenizer.py layoutdm.yaml.

Forward (corruption) process

Each token follows a mask-and-replace transition. Per step it keeps its value with probability \alpha_t, moves to a uniformly chosen other value with probability \beta_t, and moves to [MASK] with probability \gamma_t. [MASK] never leaves. The corruption is applied to each variable \bm{z}_t independently §3.1.

\mathbf{Q}_t = \begin{bmatrix} \alpha_t+\beta_t & \beta_t & \cdots & \beta_t & 0 \\ \beta_t & \alpha_t+\beta_t & \cdots & \beta_t & 0 \\ \vdots & \vdots & \ddots & \beta_t & 0 \\ \beta_t & \beta_t & \beta_t & \alpha_t+\beta_t & 0 \\ \gamma_t & \gamma_t & \gamma_t & \gamma_t & 1 \end{bmatrix} eq. 3

• Corruption is modality-wise: separate matrices \mathbf{Q}^c_t, \mathbf{Q}^x_t, \mathbf{Q}^y_t, \mathbf{Q}^w_t, \mathbf{Q}^h_t, sized (C+2)\times(C+2) for categories and (B+2)\times(B+2) for geometry §3.2. The [PAD] and [MASK] entries count toward those sizes.
• The schedules \alpha_t and \gamma_t follow VQDiffusion; the paper does not give their values §4.4. The official code's defaults are given under Hyperparameters.
• Diffusion has T = 100 steps §4.4. Sampling starts from an all-[MASK] sequence §3.1.

The posterior used in training is the standard one, obtained by Bayes' rule over the forward chain eq. 2. The paper writes it for the scalar case and then applies it per token §3.1.

q(\bm{z}_{t-1}\mid\bm{z}_t,\bm{z}_0) = \frac{q(\bm{z}_t\mid\bm{z}_{t-1})\,q(\bm{z}_{t-1}\mid\bm{z}_0)}{q(\bm{z}_t\mid\bm{z}_0)} eq. 2

Reverse process and training

• The reverse model follows D3PM: a network predicts clean-token logits, which are combined with the posterior to give the reverse distribution §3.1.

p_\theta(\bm{z}_{t-1}\mid\bm{z}_t) \propto \sum_{\tilde{\bm{z}}_0} q(\bm{z}_{t-1}\mid\bm{z}_t,\tilde{\bm{z}}_0)\,\tilde{p}_\theta(\tilde{\bm{z}}_0\mid\bm{z}_t) eq. 4

• Training minimizes a variational bound plus a weighted denoising term §3.1.

\mathcal{L}_\lambda = \mathcal{L}_\mathrm{vb} + \lambda\,\mathbb{E}\left[-\log \tilde{p}_\theta(\bm{z}_0\mid\bm{z}_t)\right] eq. 5

with \lambda = 0.1 §4.4.

• The network is a bidirectional Transformer encoder that reads the flattened sequence §3.2. Positional information is decoupled into an element index and an attribute index, each with its own learned embedding, summed per token §3.2 nn_lib.py.
• Backbone: 4 layers, 8 heads, 512 model dimensions, 2048 hidden units, dropout 0.1 §4.4. The official LayoutDM config changes this (see Paper vs code).
• Optimizer: AdamW, learning rate 5\times10^{-4}, \beta = (0.9, 0.98) §4.4.
• Epochs: 50 on Rico and 20 on PubLayNet app. A.3.

Conditional generation

The denoiser is never retrained. Conditions are imposed in every reverse step §3.3.

• Strong constraints (masking). After each step, known fields are replaced by their known values: \bm{z}_{t-1} = \bm{m}\odot\bm{z}^\text{known} + (\bm{1}-\bm{m})\odot\hat{\bm{z}}_{t-1}, where \bm{m} marks the known fields §3.3. Known fields are also fixed in the initial state \bm{z}_T §3.3.
• Weak constraints (logit adjustment). The reverse distribution is shifted in log space by a prior \bm{\pi} with weight \lambda_\pi eq. 6. The prior is either hand-written eq. 8 or the negative gradient of a sum of differentiable losses eq. 7 §3.3.

\log \hat{p}_\theta(\bm{z}_{t-1}\mid\bm{z}_t) \propto \log p_\theta(\bm{z}_{t-1}\mid\bm{z}_t) + \lambda_\pi\bm{\pi} eq. 6

\bm{\pi} = -\nabla_{p_\theta(\bm{z}_{t-1}\mid\bm{z}_t)} \sum_{i=1}^{L} \mathcal{L}_i\left(p_\theta(\bm{z}_{t-1}\mid\bm{z}_t)\right) eq. 7

\left[\bm{\pi}^i_x\right]_j = \begin{cases} 1 & \text{if } |\mathrm{loc}(j)-\hat{x}_i| < m \text{ and } j \in X \\ 0 & \text{otherwise} \end{cases} eq. 8

where \hat{x}_i is the noisy observed coordinate, m is a margin, X is the set of vocabulary indices for x, and \mathrm{loc}(j) is the centroid of bin j §3.3.

Tasks §4.3:

• Category to size and position (C→S+P) and category and size to position (C+S→P): the given fields are masked in as strong constraints.
• Completion: 0% to 20% of a complete layout's elements are kept as known §4.3.
• Refinement: the category is known and the geometry is a noisy copy of a real layout, with Gaussian noise of variance 0.01 §4.3. The geometry prior is the hand-written rule eq. 8, and the categories are fixed by masking §4.3.
• Relationship: the category is known and about 10% of element pairs have size or location relations §4.3. The loss in app. A.2 penalizes violated relations. Because the bins are discrete, the paper converts each probability vector into an expected coordinate \bar{x} = \sum_n p(x=n)\,\mathrm{loc}(n) §4.5. The adjustment is applied three times per diffusion step §4.5.
• Unconditional: no conditions §4.3.

Fast sampling. Sampling with a step size \Delta uses T/\Delta denoising steps §4.6. The paper gives a modified reverse distribution for this §4.6; it is not reproduced here (see Paper vs code).

Algorithms

Training, one batch. Element shuffling and the variational term follow the paper; the auxiliary weight is the paper's constant (see Paper vs code for the code's version).

# Training step on one batch of layouts  [§3.1] [eq. 5]
# Input:  layouts: (B, 25) element ids, padded with [PAD]
# Output: scalar loss

def training_step(layouts):
    # Shuffle the elements of each layout at random  [§3.2]
    tokens = flatten_shuffled(layouts)  # (B, 125)
    # Draw one timestep per example; the paper does not state the law
    t = sample_timesteps(batch_size)  # (B,), 0 <= t < T
    # Corrupt each attribute with its own matrix  [§3.2] [eq. 3]
    noisy = corrupt_per_attribute(tokens, t)
    # Predict clean-token logits from the noisy sequence  [§3.1] [eq. 4]
    logits = denoiser(noisy, t)  # (B, 125, K)
    # Variational term; the decoder term applies at t = 0  [§3.1] [eq. 5]
    loss_vb = variational_term(tokens, noisy, logits, t)
    # Denoising term with the constant weight lambda = 0.1  [§4.4] [eq. 5]
    loss_aux = 0.1 * cross_entropy(logits, tokens)
    return loss_vb + loss_aux

Conditional sampling. Known fields are fixed at every step, and the weak prior is added in log space.

# Constrained sampling  [§3.3] [eq. 4] [eq. 6]
# Input:  known (B, 125) ids, known_mask (B, 125) bool, optional prior
# Output: sampled layouts (B, 125) ids

def sample_layout(known, known_mask, prior=None, T=100):
    # Start from [MASK] everywhere; known fields are fixed  [§3.2] [§3.3]
    state = where(known_mask, known, MASK_ID)
    for t in reversed(range(T)):
        # Reverse distribution from the denoiser  [eq. 4]
        log_probs = log_posterior(denoiser(state, t), state, t)
        # Add the weak prior in log space; the code adds it only for
        # t >= 10 (see Paper vs code)  [eq. 6] [eq. 7] [eq. 8]
        if prior is not None:
            log_probs = log_probs + lambda_pi * prior(log_probs)
        # Re-impose the known fields after every step  [§3.3]
        log_probs = where(
            known_mask, known_log_probs(known), log_probs
        )
        state = sample(log_probs)
    return state

Hyperparameters

Results

Author-reported. The paper averages three independent trials per model §4.4, and each table prints a spread value after the mean. Hardware is not stated. For Max., Sim. and Align., the paper does not state the unit of the ± values; they are reproduced as printed.

Table 1: conditional generation given partially known fields tab. 1. Columns are FID ↓ and Max. ↑ for Rico and PubLayNet, per task. † marks BLT trained with [PAD] as an extra vocabulary, because the original model cannot do unordered completion in practice tab. 1. Real data is reported for reference.

PLN = PubLayNet. The paper's bold and underline markings are not reproduced. The paper says LayoutDM outperforms all but LayoutTrans on completion §4.5.

Table 2: unconditional generation tab. 2. Align. is alignment (lower is better in the paper's table); Real data is the validation set.

The paper notes LayoutDM has the best FID among random-order models; the fixed-order LayoutTrans is better on Rico §4.5. Its own Rico FID is second only to LayoutTrans-fixed.

Table 3: refinement tab. 3. Sim. is DocSim similarity to the ground truth.

Table 4: ablation of layout-specific components on Rico, unconditional tab. 4.

Other results, in words and in appendix tables.

• Relationship constraints fig. 5: author-reported that LayoutDM is comparable to NDN-partial on Rico and well ahead of it on PubLayNet, and behind CLG-LO on both. The average runtime is 4.0 s for CLG-LO and 0.5 s for LayoutDM §4.5.
• Speed-quality trade-off for C+S→P fig. 6: the paper reports LayoutDM as the best trade-off among task-agnostic models, except the task-specific LayoutGAN++, which runs under 10 ms §4.6.
• State space app. B: discrete LayoutDM has FID 6.65 at 100 steps; continuous DiffusionLM reaches 22.8 at best (1000 DDPM steps) app. B.
• Alignment and overlap app. B: fixed-length tasks are comparable to VQDiffusion* and BART; for variable-length tasks, BART and LayoutTrans are moderately better than LayoutDM.

Paper vs code

Official code at 873b5eeb, which is the commit this card is based on. The code is Apache-2.0 and the authors' lab owns it. Discrepancies, each with both citations:

1. Dropout. Paper: 0.1 §4.4. The LayoutDM experiment config sets encoder-layer dropout to 0.0 layoutdm.yaml. The Transformer's embedding dropout is hard-coded to 0.1 nn_lib.py.
2. Backbone width. Paper: 512 dimensions and 2048 hidden units §4.4. LayoutDM shrinks both by 29/32, giving 464 and 1856, "for fair comparison" layoutdm.py util.py.
3. Epochs on PubLayNet. Paper: 20 app. A.3. The code default is 50 hydra_configs.py, and no dataset config overrides it.
4. Relation prior timing. Paper: the adjustment is applied three times to each diffusion step §4.5. Code: it is applied only when t \ge 10 logit_adjustment.py.
5. Relation prior scale. Paper: \bm{\pi} = -\nabla\sum_i\mathcal{L}_i eq. 7. Code averages the constraint losses logit_adjustment.py and takes one SGD step with learning rate relation_lambda = 3e6 logit_adjustment.py. The paper's \lambda_\pi for relations is not stated.
6. Auxiliary loss weight. Paper: a constant \lambda = 0.1 eq. 5. Code scales the auxiliary term per step by (1 - t/T) + 1, because adaptive_auxiliary_loss defaults to true constrained.py base.py.
7. Nucleus sampling default. Paper: p is searched over {0.90, 0.95, 0.99, 1.0} app. A.3. The test config defaults to random sampling, and top-p is applied only when sampling=top_p hydra_configs.py test.py.
8. FID network. Paper: an improved layout feature extractor from Kikuchi et al. §4.2. Code: the FID model is FIDNetV3, loaded from fid_weight_dir model.py main.py. The card did not verify that these are the same network.
9. Fast sampling. Paper: a modified reverse distribution over \Delta-step jumps §4.6. Code: a skip offset inside the posterior call base.py. The card did not verify the equivalence.
10. Refinement noise. Paper: variance 0.01 §4.3. Code: standard deviation 0.1, which is the same variance task.py. Consistent.

Code-only details the paper does not state: the LR schedule (see Hyperparameters), gradient clipping at 1.0, weight decay 0.1, batch size 64, importance sampling of timesteps once every step has more than 10 loss records base.py, and the timestep conditioning via adaptive LayerNorm layoutdm.yaml. Also code-only is forcing [PAD] off inside elements already present, in conditional tasks base.py.

Pitfalls

• The default q_type is invalid. LayoutDM.__init__ defaults to q_type="single", which fails the Q_TYPES assertion. Use "constrained" for the paper's modality-wise model, which the experiment config does layoutdm.py layoutdm.py. "vanilla" is a single matrix over all tokens. The card did not check which configuration produced the ablation in Table 4.
• Warmup does nothing in the official config. The warmup LR is 5e-4 and the base LR is 5e-4 reduce_lr_on_plateau_with_warmup.yaml layoutdm.yaml. The warmup step is therefore zero scheduler.py. Only the plateau rule changes the LR after warmup.
• \beta_t is derived, not free. Column sums of \mathbf{Q}_t must equal 1, which gives \beta_t = (1-\alpha_t-\gamma_t)/N with N the number of non-[MASK] classes constrained.py. This derivation is ours and the code follows the same rule. The paper does not state it.
• Conditioning uses the noisy input as the prior centre. For refinement the prior is built from the noisy bins, not the clean ones task.py. The prior is only added to fields that are not known task.py.
• Refinement noise is redrawn on every call. get_cond draws new Gaussian noise each time it runs task.py. Seed the run to reproduce a table.
• Layouts with more than 25 elements are dropped, so the evaluated distribution is not the raw dataset §4.1 rico.py.
• Results are not comparable across papers. The paper says the datasets lack official splits and FID networks differ across works §4.4, footnote. Compare only within this paper.
• Std values for Max., Sim. and Align. The paper does not give their unit; read them as printed.
• Model selection uses the best validation loss main.py. The paper also searches hyperparameters on the validation set §4.4.
• Repository status. The GitHub API reports the repository as archived.
• Supplementary equations are not numbered in the main sequence. Cite them by section, such as app. A.2, not by number.

Lineage

• Builds on: D3PM, the discrete state-space diffusion framework §3.1; VQDiffusion's mask-and-replace transition §3.1; the Transformer encoder §3.2. Neither has a card on this site.
• Compared against, task-agnostic: LayoutTrans, BLT, BART, MaskGIT*, VQDiffusion* §4.3.
• Compared against, task-specific: LayoutVAE, NDN-none, LayoutGAN++ and CLG-LO, RUITE §4.3 §4.5.
• Superseded by: not known from wiki content.