Two-Level Query Mechanism in Transformers

• Two-Level Query Mechanism is a hierarchical approach that uses room-level and corner-level queries to generate structured outputs without manual post-processing.
• It employs transformer decoders with multi-head self and cross-attention for iterative refinement of geometric features in tasks like floorplan reconstruction.
• Advantages include native variable-length handling, end-to-end training via Hungarian matching, and state-of-the-art performance on benchmarks such as Structured3D and SceneCAD.

A two-level query mechanism is a structural approach in transformer-based neural architectures tailored for tasks requiring the generation of hierarchical and variable-length outputs, such as multi-room floorplan reconstruction or scene graph synthesis. It enables the model to emit collections of structured objects (e.g., rooms), each comprising an ordered sequence of sub-objects (e.g., polygon corners), supporting variable cardinality and length at both hierarchy levels without resorting to manual post-processing or multi-stage design. This paradigm has been instrumental in advancing state-of-the-art results in geometric scene understanding and vectorized layout modeling, particularly by RoomFormer (Yue et al., 2022).

1. Structural Principles of the Two-Level Query Mechanism

The two-level query mechanism distinguishes itself by encoding output as a set of sequences, concretely representing a floorplan as a variable-size set of room polygons, each modeled as an ordered variable-length sequence of corner vertices. In RoomFormer (Yue et al., 2022):

• Top-Level Queries ($M$): Each room is allotted a “slot” via a room-level query (maximum $M$, e.g., $M=20$). These slots function as unordered members of a prediction set.
• Second-Level Queries ($N$): For each room, a sequence of $N$ queries (usually $N=40$) models its corners, with a validity flag indicating which are active (real) corners and which are padding.
• Combined Input: The transformer decoder receives a tensor of shape $M \times N \times 2$ (room slots, corner slots, and coordinate/channel dimensions).

This structure allows hierarchical outputs: unordered at the set level (rooms), ordered at the sequence level (corners), handling both variable set cardinality and sequence length within a single pass.

2. Implementation in Transformer Architectures

Within RoomFormer (Yue et al., 2022), the two-level query mechanism is realized in the decoder, which iteratively refines both room and corner queries using multi-head self-attention and multi-scale deformable cross-attention:

• Self-Attention: All corner queries (across all rooms) attend to each other, permitting global geometric context propagation.
• Cross-Attention: Room and corner queries attend to multi-scale features from the encoded input (e.g., density images from projected point clouds).
• Iterative Refinement: After each transformer layer, corner coordinates are updated via predicted offsets. Validity predictions permit dynamic output truncation at either hierarchy level.

The separation of queries at two levels allows direct structured prediction without imposing sequence or set limits via fixed post-processing rules.

3. Polygon Matching and End-to-End Training

A core technical hurdle in variable-sized set prediction is one-to-one assignment between predicted and ground-truth polygons (rooms):

• Bipartite Matching: Padding to maximum size ($M$ for rooms, $N$ for corners) allows use of the Hungarian algorithm based on a cost function combining classification (corner validity) and coordinate errors.
• Cost Function ($\mathcal{D}$): The matching cost between predicted room $V_m$ and target $\hat{V}_{\sigma(m)}$ is

$$\mathcal{D}(V_m, \hat{V}_{\sigma(m)}) = \lambda_{cls} \sum_n\|c^m_n - \hat{c}_n^{\sigma(m)}\| + \lambda_{coord} d(P_m, \hat{P}_{\sigma(m)}),$$

where $d(\cdot,\cdot)$ computes minimal $L_1$-distance over valid cyclic permutations, handling the lack of canonical starting vertex for closed polygons.

• Losses: The final loss aggregates validity (binary cross-entropy), coordinate regression ($L_1$), and auxiliary rasterized mask loss (Dice loss using a differentiable rasterizer).

End-to-end training is thus feasible, directly optimizing structured output accuracy without recursive post-processing or iterative polygon search.

4. Advantages Over Single-Level and Heuristic Methods

The two-level query mechanism provides several distinct advantages:

Dimension	Two-level Queries	Single-level/Heuristic
Hierarchical Output	Set of ordered sequences	Flat sequence or set
Variable Size Handling	Native (slot validity flags)	Often rigid or post-processed
Training End-to-End	Yes (gradient-based matching)	Often staged or post-hoc matching
Scalability	Fast (0.01s/sample (Yue et al., 2022))	Slower (e.g., >0.1–1s for HEAT)
Generalization	Robust to new/external data	Prone to error-propagation

Two-level queries also outperform flat or single-level queries in ablations, and permit seamless extension to predict auxiliary elements (e.g., doors, windows, semantic room types) through decoder variants.

5. Extension to Semantic and Architectural Elements

RoomFormer’s two-level paradigm generalizes to multitask prediction scenarios:

• Room Type Classification: Aggregates corner-level features per detected room, applies a linear head and softmax, permitting direct semantic enrichment.
• Architectural Elements (Doors/Windows): Predict lines as degenerate polygons (2-corner sequences), realized via additional decoder heads—either sharing the two-level structure or using a specialized single-level line decoder.

This modularity facilitates simultaneous geometric and semantic output within a unified transformer design.

6. Empirical Performance and Impact

Empirical results on Structured3D and SceneCAD datasets demonstrate the efficacy of two-level queries for floorplan reconstruction (Yue et al., 2022), with state-of-the-art metrics:

• Room F1: 97.3%
• Corner F1: 87.2%
• Angle F1: 81.2%
• Runtime: 0.01s/sample (over 10× faster than previous pipelines)
• Generalization: RoomFormer, trained on Structured3D, widely outperforms previous SOTA when evaluated on SceneCAD (IoU 74.0 vs 52.5 for HEAT).

A plausible implication is that two-level queries offer a foundational advance for any structured scene synthesis task requiring hierarchical, variable-cardinality output, enabling fast, robust, and end-to-end trainable models without recourse to custom procedural post-processing. The mechanism has been adopted and adapted in subsequent work, confirming its general utility for vectorized geometric modeling.