LayoutFlow: Flow Matching for
Layout Generation

Over the years, layout generation has been investigated using various methods. Early works started exploring generating layouts by minimizing an energy function based on pre-defined constraints [37, 36]. With the advent of generative machine learning techniques, however, research shifted to a more data-centric learning approach and started to use Variational Autoencoders (VAE) [47, 13, 2, 27], and Generative Adversarial Networks (GAN) [29, 30, 25]. While earlier works explored a variety of different architectures, such as Graph Convolutional Networks [27], Long-Short Term Memories (LSTM) [13, 24] or Transformers [47, 45, 14], most recent research has converged to using a Transformer-based architecture due to its flexibility and strong performance.

Diffusion-based Models. As diffusion models have proven to generate more diverse data and are generally easy to train, especially compared to GANs, most recent research has shifted towards exploring the diffusion process for layout generation tasks. Since layouts are represented by categorical data defining types of elements and continuous numerical values describing the element position and size, applying standard diffusion models is not straightforward. As a result, different approaches have incorporated a diffusion loss into their training process, which can be divided into continuous [18] and discrete diffusion models [3]. For the discrete case, the continuous coordinates are quantized and interpreted as different states along with the categorical element type [22, 20, 48, 15]. In this discrete scenario, the forward diffusion process is modeled as a random walk between discrete states that, in addition to the quantized states, also include a mask state, which removes elements from the canvas. While LayoutDM [22] and LDGM [20] rely on the mask-and-replace strategy proposed in [12], LayoutDiffusion [48] introduces a new mild forward process that is closer to the continuous process while still increasingly masking elements over time.

On the other hand, continuous diffusion models have been less common as they have not attained the same performance as discrete models. In particular, continuous models have struggled to produce layouts with well-aligned elements. The first continuous layout generation model based on diffusion was proposed by Chai et al. [4], but is restricted to generating continuous element coordinates given categorical type data as a condition. Cheng et al. [8] introduced a latent diffusion model for layout design focusing on user constraints by embedding the layout into continuous representation, on which the diffusion process is performed. The first continuous diffusion-based layout generation model acting on the layout representation and able to handle unconditional generation was proposed by Chen et al. [5]. Their work addresses the alignment issue observed for continuous diffusion by introducing a specific regularization loss. Lastly, Levi et al. [28] proposed DLT, which applies a continuous diffusion process on the element coordinates and a discrete process on the element types.

Large Language Models. Motivated by the strong generalization capabilities of Large Language Models (LLM), the most recent works [31, 43, 11] have investigated the layout generation abilities of LLMs. While LLMs provide some interesting applications, such as zero-shot synthesis, they typically consist of billions of parameters, resulting in long inference times. Moreover, current models are limited to conditional layout generation. In contrast, our proposed approach aims to offer a lightweight model for unconditional and conditional layout generation.

Flow-based models are based on learning a mapping between samples from a simpler distribution to a data distribution. Over the years, there have been various methods, such as Normalizing Flows or Continuous Normalizing Flows, which have used neural networks to estimate such a mapping, which can subsequently be used to sample from the learned distribution [26]. Until recently, the biggest issue with flow-based models has been the need to backpropagate through an ODE during training. However, recent models have proposed training algorithms that only require explicitly solving an ODE during inference [44, 32, 1, 33, 35] demonstrating impressive results and outperforming diffusion models on the image generation task [34]. Our proposed model builds on this improvement by utilizing Flow Matching for layout generation while offering an improved geometrical interpretation.

Diffusion models, which were first proposed by Sohl-Dickenstein et al. [40] and later popularized by Ho et al. [18], can be formulated similarly to flow-based models. In fact, the term Flow Matching was inspired by the similar Score Matching loss [21] used by diffusion models. However, flows are characterized by an ODE, whereas the diffusion process can be formulated as an SDE with an additional stochastic component described by Brownian Motion [42]. As a result, diffusion models are restricted to a Gaussian as a base distribution, while flow-based models allow for more flexibility, which we also explore in our experiments. Notably, Song et al. [41] proposed DDIM, an alternative sampling method that makes it possible to sample using an ODE formulation on a model trained with an SDE forward process, effectively leading to smoother trajectories. In the context of layout generation, however, models have been using the SDE formulation following DDPM [18]. Our argument for introducing Flow Matching instead of just using DDIM for smoother trajectories lies in its increased flexibility and simplicity. Moreover, if the goal is to sample using an ODE, it seems more natural to also use an ODE during training instead of the more complex SDE used in diffusion.

A layout typically consists of various elements, such as text, tables, or images placed on a canvas. Mathematically, we can describe a layout $\mathcal{A}$ as a set of $N$ elements $\mathbf{e}^{k}$ with $k\in[1,N]$ in the following way

where we split each element into its geometrical information $\mathbf{g}^{k}=(c_{x}^{k},c_{y}^{k},w^{k},h^{k})$ described by the bounding box center $(c^{k}_{x},c^{k}_{y})$, width $w^{k}$ and height $h^{k}$, and its category information $\mathbf{a}^{k}$ describing the element type, such as text or image. In our considerations, the elements on the canvas do not possess an inherent order, which is why the canvas is considered a set of elements. Layout generation can be divided into various tasks depending on the conditioning assumptions. For example, in the unconditional generation setting, it is assumed that no prior knowledge is available, while for the refinement task, a complete but coarse layout is given. Since Flow-Matching requires the input data to be continuous, we convert the categorical information $\mathbf{a}^{k}$ into a continuous embedding, denoted as $\tilde{\mathbf{a}}^{k}$, using $B$ Analog Bits [7]. To learn the flow, we define a data sample $\mathbf{x}\in\mathbb{R}^{(4+B)\cdot N_{max}}$ as the concatenation of all the elements $\mathbf{g}^{k}$ and $\tilde{\mathbf{a}}^{k}$ in a single layout. If the number of elements is smaller than $N_{max}$, the maximum number of elements contained in a single layout across the entire dataset, we pad the remaining dimensions.

To generate layouts, we want to find a way to sample from a data distribution $p_{\text{data}}(\mathbf{x})$, whereas we only have access to a limited amount of samples $\mathbf{x}$ through our training data. More specifically, our goal is to train a neural network that allows us to sample from the underlying data distribution and guide the sampling based on certain conditional constraints. To that end, we propose LayoutFlow, a layout generation model based on Flow Matching.

An overview of the training procedure for the type-conditioned scenario is illustrated in Fig. 2 and summarized in Algorithm 1 for the unconditional case. In the first step, we randomly sample a layout $\mathbf{x}_{0}$ from a prior distribution $p_{0}$, e.g., a Gaussian distribution. This sample will typically not be well-aligned, and some elements might not be located on the canvas. In addition, we also randomly select a time $t$, which determines the intermediate sample $\mathbf{x}_{t}$ on the flow trajectory used during training. For LayoutFlow, we train on a simple linear trajectory following [44, 32]. Therefore, the intermediate sample is simply obtained by linearly interpolating between the ground truth layout $\mathbf{x}_{1}$ and the initial sample $\mathbf{x}_{0}$. The same trajectory is not used again during training since $\mathbf{x}_{0}$ and $t$ change with every iteration.

The network $f_{\theta}$ is trained to output the conditional vector field, corresponding to the derivative of the trajectory as defined in Eq. 1. This derivative is simply a constant for the linear trajectory case, specifically the difference between $\mathbf{x}_{0}$ and $\mathbf{x}_{1}$. Intuitively, the network learns to output a direction pointing toward a data sample, which can be used during sampling to move the initial data samples toward a better prediction. While we train our network using linear trajectories, this does not mean that the model will also produce straight trajectories during sampling. As mentioned in Eq. 4, the linear trajectory, which is derived by the conditional vector field in Eq. 4, acts as a proxy to learn the actual vector field in Eq. 3 defined by Flow Matching.

In our experiments, we found that using only the Mean Squared Error (MSE) as a loss function tends to lead to poorly aligned layouts, which results in suboptimal perceptual quality. This observation is similar to the issue described in [5]. We hypothesize that this can be attributed to the fact that the MSE penalizes all mistakes within the same distance equally, regardless of direction, whereas the alignment between elements is usually only measured along the horizontal and vertical direction. As a result, the minimization objective does not fully align with the perceptual quality. To tackle this issue, we add an $L_{1}$ regularization on the geometrical part of the output $u^{\mathbf{g}}_{\theta}$ as the $L_{1}$ loss encourages sparsity of the error vector, which can also be interpreted as minimizing the error along the axis dimensions. Therefore, our final training loss can be written as

where $\lambda$ is a hyperparameter.

After training our neural network to predict a vector field, we can generate new layouts by sampling from our initial distribution and solving the ODE describing the flow as defined in Eq. 1. In practice, this can be done using any numerical ODE solver. Figure 4 illustrates the sampling process using the Euler method [10] to solve the ODE. For a fixed number of evaluation steps $T$, the initial sample $x_{0}$ is moved along the direction predicted by the network in an autoregressive manner. Therefore, the backward process can be denoted as

with $i\in\mathbb{N}$ and $i\in[0,T-1]$ indicating how far along the sample has been moved.

Here, we describe the actual implementation of the network $f_{\theta}$ to efficiently handle the layout. We first embed the elements of the sampled intermediate layout $\mathbf{x}_{t}$ and then use a Transformer to predict the conditional vector field as depicted in Fig. 2. First, we separately embed the geometrical information $\mathbf{g}^{k}_{t}$ and the element type $\tilde{\mathbf{a}}^{k}_{t}$ using a linear layer. A conditioning mask is embedded and added to the geometrical and type embedding to inform the network about which parts of the input act as a condition. The example in Fig. 2 illustrates this for the type-conditioned case, where the network is informed through the masks $\mathbf{m}_{type}$ and $\mathbf{m}_{pos}$ that $\tilde{\mathbf{a}}^{k}$ is taken from $\mathbf{x}_{1}$ while only the geometrical information $\mathbf{g}^{k}_{t}$ was taken from the training trajectory. In a subsequent step, we concatenate both embeddings and pass them through another linear layer to obtain a high-dimensional embedding $\mathbf{f}_{t}^{k}$ representing an element of the intermediate layout $\mathbf{x}_{t}$. Next, all element embeddings $\mathbf{f}_{t}^{k}$ are processed by a transformer encoder network together with the temporal information $t$ incorporated in each layer. Finally, the output of the Transformer is projected back to the same dimension as the input layout using a linear layer representing the vector field.

As layout generation involves various tasks with different assumptions about available knowledge, we employ a conditioning mechanism to perform these tasks with just a single model. Similarly to [5, 20, 28], we mark the conditioning input by including the information in the element embedding while distinguishing between the four different scenarios depicted in Fig. 4. During training, we simply randomly sample each mask with the same probability. Alternatively, it is also possible to only train an unconditional model and impose the conditions during sampling through a guidance strategy, similar to [22, 48]. However, we empirically found that including conditioning during training works better. For a more detailed analysis, we refer to the supplementary material.

Datasets. We evaluate the performance of our network on the RICO [9] and PubLayNet [49] datasets, following previous methods. RICO consists of over 66k User Interface (UI) layouts with 25 element categories, while PubLayNet comprises over 360k document layouts annotated with 5 different element types. We train our network using the dataset split described in [48, 23], which discards layouts containing more than 20 elements.

Implementation Details. Similar to LayoutDM and DLT, we employ a 4-layer Transformer network with 8 attention heads and a model dimension of 512 for a fair comparison. The model is trained using the AdamW optimizer with a learning rate of 0.0005. In addition to the MSE Flow Matching loss, we regularize the training by calculating the $L_{1}$ loss between the geometrical information as described in Eq. 6 with $\lambda=0.2$. For inference, we obtain the flow using a NeuralODE solver [6] for 100 steps with the Euler method.

Evaluation Metrics. Following previous work [48, 4, 28, 5], we evaluate our method as well as comparable approaches using Frechet Inception Distance (FID) [16], Maximum Intersection over Union (mIoU) [25], Alignment [27], and Overlap [29]. FID and mIoU inherently evaluate both fidelity and diversity [17] of the generated results. For the FID calculation, we use the same network with identical weights as LayoutDiffusion [48] to measure the similarity between the generated layouts and the original dataset based on the feature space. The mIoU metric compares sets of layouts and measures their similarity by optimally matching generated and real layouts that maximize the average IoU. On the other hand, Alignment and Overlap capture the fidelity property. Nonetheless, both metrics should be judged with respect to a reference dataset, where the best outputs lead to similar values.

Comparison with Existing Approaches. To validate the performance of our proposed method, we compare it to state-of-the-art techniques. Since there is no standard way of splitting the datasets, we need to retrain the models that do not follow the dataset split used in LayoutDiffusion [48] for a fair comparison. This fact limits us to only compare to methods that have made their code publicly available. Therefore, we compare our method against the two discrete diffusion models LayoutDM [4] and LayoutDiffusion [48] as well as the discrete-continuous diffusion model DLT [28]. Where applicable, we additionally compare LayoutFlow to the non-diffusion-based models LayoutTransformer [14], LayoutFormer++ [23], NDN-none [27] and RUITE [38].

Tasks. We test LayoutFlow and existing approaches on various common unconditional and conditional layout generation tasks. Un-Gen describes the layout generation task without any constraints. For the conditional scenario, we consider Gen-Type as assuming the element types as given, while Gen-TypeSize assumes element types and sizes are known. Lastly, we consider the Completion task, which takes in a partial layout that is missing up to 20% of its elements, and the Refinement task following [48] with a standard deviation of 0.01.

Quantitative Evaluation. In Tab. 10, we report the results of our experiments comparing LayoutFlow with existing approaches for different layout generation tasks. In terms of FID, our model outperforms state-of-the-art methods across all three tasks, except for a close second place on the PubLayNet Un-Gen task, proving its strong capabilities to model the underlying sample distribution. In particular, the FID score of 2.37 produced by LayoutFlow on the unconditional generation for RICO almost matches the score obtained when comparing the validation with the test dataset. Note that for the conditional tasks, the FID score can become even lower than the validation dataset since the generated images share more similarities through conditioning. Regarding the geometrical metrics Alignment, Overlap, and mIoU, our method consistently outperforms LayoutDM and DLT while usually producing values close to LayoutDiffusion, a significantly larger model with 85M parameters compared to around 15M used by LayoutFlow. A comparison between the quality and the sampling speed is illustrated in Fig. 8, clearly showing that our method closes the gap between inference speed and performance. While non-diffusion models outperform diffusion models on some geometrical metrics, the drastically higher FID score implies that the generated results lack diversity.

On the completion task, shown in Tab. 2, LayoutFlow is able to clearly outperform LayoutDM, the only diffusion model that can handle that task, except for falling behind on Alignment for RICO. On the refinement task, LayoutFlow significantly improves the noisy layouts on PubLayNet, while being the second-best model for RICO, underlining its broad capabilities.

Qualitative Evaluation. We provide some qualitative examples in Fig. 6 for the unconditional Fig. 7 as well as for the conditional generation task on the RICO dataset. More samples for the other tasks and PubLayNet can be found in the supplementary material. Overall, LayoutFlow shows a strong performance across all tasks, producing visually pleasing results that resemble real layouts.

Limitations. We discuss the limitations of our proposed methods in the supplementary material.

In order to justify our design choices, we perform an ablation study on unconditional layout generation using the RICO dataset.

Diffusion vs. Flow. We show that employing Flow Matching as the generative model is essential by also training the model used for LayoutFlow and only substituting Flow Matching with diffusion. We test two methods for sampling from the diffusion-based model, i.e., DDPM and DDIM. As shown in Tab. 4, neither approach is able to reach the same performance as LayoutFlow, and both seem to be particularly lacking in terms of Overlap. Interestingly, DDIM outperforms DDPM by a large margin, which might be related to DDIM sampling based on an ODE instead of an SDE.

Regularization Loss. We validate the efficacy of our regularization loss proposed in Eq. 6 and the choice of $\lambda$ in Tab. 4. It can be clearly observed that a stronger regularization using the $L1$-loss provides better alignment but increases the FID. Therefore, we considered the trade-off across all metrics and chose $\lambda=0.2$ for LayoutFlow.

Initial Distribution. Since flow-based models allow for arbitrary initial distributions, we explore two alternatives. First, we try a uniform distribution to ensure all initial samples are placed on the canvas. In addition, as Analog Bits have a smaller sample space, we tested using a Gaussian for the geometrical elements and a uniform distribution for the Analog Bits. Altogether, the results in Tab. 6 indicate that a Gaussian initial distribution provides the best overall performance measured by the FID score, whereas a uniform distribution provides stronger results in terms of geometrical metrics, such as Alignment and Overlap.

Training Trajectories. As long as the conditional vector field fulfills the criteria presented in Sec. 3, it can be freely chosen for training. In addition to the linear training trajectory proposed in [32, 44], we also explore the sine/cosine interpolation introduced in [1] and a sine-based interpolation, as shown in Tab. 6. The choice of training trajectories only slightly affects the performance.