Autoregressive Mosaics: Probing 2D Spatial Reasoning in Text-Only Language Models

The paper introduces AM-Bench, a novel benchmark designed to measure allocentric 2D layout at mosaic granularity. This framework serves as a necessary-condition probe for embodied spatial competence by evaluating how well text-only language models can perform complex visual reasoning tasks. The study is significant because it systematically tests whether modern generative AI models possess inherent spatial planning capabilities, moving beyond simple linguistic understanding to assess structured, geometric comprehension.

Scope and Benchmark Design

AM-Bench measures 2D layout by requiring the model to generate autoregressive mosaics, which are structural compositions of shapes. The benchmark is designed to be comprehensive, covering various levels of spatial difficulty. The test suite includes multiple types of challenges:

1. Iconic Prompts: These prompts require recognizing and reproducing complex visual structures, such as A traffic light (tall dark rectangle, th...) or A yellow sun with eight radiating rays o...

2. Periodicity/Tiling: The benchmark specifically tests for periodicity—repeating grids or tilings—an axis that was historically difficult to build into language models.

3. Compositional Inversion: This task assesses the model's ability to handle complex structural relationships, such as an Iconic/Compositional inversion.

Evaluation Metrics and Analysis

The evaluation utilizes multiple metrics to quantify spatial performance, including the J1 score (a measure of layout quality) and attention standard deviation. The analysis also employs token-generation entropy to gauge consistency. For instance, Figure 16 illustrates the relationship between entropy and local-structure distance across eight different models, showing a strong correlation (Spearman rho = 0.93). Performance is measured against various model architectures and sizes, including Qwen, Gemma, GLM, and Llama/CodeLlama families (8B–34B).

Observed Model Behaviors and Limitations

The study reports distinct behavioral patterns across the tested models. For example, in the DINO attention maps (Figure 17), the top model (GLM-4 32B) often demonstrates highly focused attention on key structural elements, whereas the worst model (CodeLlama 34B) shows less precise localization. However, researchers caution that these findings are not definitive: it is still not a controlled scaling-law sweep, and whether the observed patterns hold at larger scales or for closed-weight models remains untested.

Future Directions and Technical Depth

To further validate its claims, the research proposes several avenues for expansion. These include:

• Scaling the benchmark to encompass broader spatial categories through a larger multiannotator human validation run.

• Introducing other output mediums as an additional code-prior ablation to investigate the effect of the output medium.

• Applying advanced causal analysis techniques, specifically activation patching for a direct, causal analysis of the models’ internal spatial planning, to substantiate behavioral claims at the representation level.

The authors also provide extensive technical detail regarding data quality, noting that some Layout samples are missing a second-judge score and that the pooled foreground mask can be empty for approximately 3% of structure-coherence cells when describing several small, spatially diffuse elements.

The provided text excerpts detail sophisticated, multi-faceted evaluations of current LLMs, focusing heavily on spatial reasoning (Layout/Iconic prompts), code generation reliability (Entropy analysis), and multimodal understanding (DINO attention). The gaps highlighted—especially regarding periodicity, metric continuity, and the limited scope of the current benchmarks—represent critical areas for immediate architectural and methodological improvement.

Here are the specific, high-impact improvements required to advance AI systems using this research foundation.

Improvement 1: Implementing Causal Activation Patching for Internal Spatial Planning.

• Technical Action: Do not rely solely on surface-level output metrics (like Jaccard Index or layout scores). Integrate and utilize activation patching techniques during inference to perform a direct, causal analysis of the model's internal spatial representations. Specifically, target the transformer layers responsible for compositionality and object permanence within the visual embedding space.

• Rationale: The paper notes that behavioral claims should be substantiated at the representation level. We must move beyond correlation to causation when assessing spatial reasoning failures (e.g., why does it fail on periodicity?).

• Improved System Capability: The AI system can generate an Explainable Spatial Failure Report. When a layout task fails, the system doesn't just output Wrong. It identifies which specific internal attention head or activation subspace failed to maintain object coherence (e.g., Failure point: Loss of relative depth metric between object A and background B in Layer 12). This allows for targeted architectural retraining.

Improvement 2: Integrating a Metric-Aware, Multiscale Representation Module.

• Technical Action: Modify the backbone architecture to explicitly model metric egocentric continuity, depth, and occlusion as first-class citizens of the latent space. This requires coupling the standard ViT encoder with a dedicated depth estimation module (e.g., using a monocular depth prediction network) that is jointly trained with the layout task loss function (L total = L layout + lambda 1 L depth + lambda 2 L metric continuity).

• Rationale: The current benchmark scope is too narrow (does not cover metric egocentric continuity, depth, or occlusion). A robust AI must understand 3D spatial relationships from 2D prompts.

• Improved System Capability: The AI can perform True Embodied Spatial Planning. Given a prompt like "Place the chair behind the table and slightly above the potted plant," it generates not just a plausible image, but an internal, consistent 3D coordinate map that respects relative depth and occlusion rules, making it suitable for robotic control or advanced simulation environments.

Improvement 3: Developing and Mandating the Periodicity/Tiling Stress-Test Ladder.

• Technical Action: Formalize the novel periodicity rung into a mandatory, scaled component of all spatial evaluation suites. This set must include tasks that require understanding repeating patterns, rotational symmetries, and tiling constraints (e.g., tessellations with specific edge matching rules).

• Rationale: The paper explicitly notes that periodicity is the most likely axis of degradation and remains untested. Failure here indicates a fundamental weakness in pattern recognition beyond simple compositionality.

• Improved System Capability: The AI can demonstrate Advanced Pattern Synthesis and Prediction. It moves beyond merely reproducing a tiled image to generating the missing component or predicting the next state in a repeating, constrained pattern, proving deep understanding of geometric rules rather than rote memorization.

Improvement 4: Implementing a Comprehensive Noise-Aware Target Set Scoring System.

• Technical Action: Formalize the handling of ambiguous target inputs (like the 17 targets needing simplification). The evaluation pipeline must incorporate an automated pre-processing step that scores the ambiguity of the target description (e.g., using graph theory metrics on touching/overlapping shapes) and penalizes models based on their failure to generalize across varying levels of noise/ambiguity.

• Rationale: Source noise is a documented weakness (0.96 vs 1.00). A production-grade system cannot fail simply because the prompt was slightly vague or contained minor overlaps that required human simplification.

• Improved System Capability: The AI achieves Robust Zero-Shot Interpretation. It can process and generate accurate output even when the input prompts are imperfect, noisy, or require implicit merging of adjacent elements (merging touching, same-color shapes), making it reliable in real-world data streams.

Improvement 5: Integrating Code Generation Entropy as a Pre-Filter and Fine-Tuning Loss.

• Technical Action: The entropy measurement must transition from being merely an analysis tool to a direct training objective. Implement a loss function component (L entropy) during fine-tuning that penalizes the model when its generated code deviates significantly from the expected low-entropy, high-consistency pattern observed in valid samples.

• Rationale: The entropy trend (Figure 16) shows a strong correlation between local structure distance and generation consistency. Using this as a loss term forces the model to learn predictable internal representations for structured data (like code or geometric coordinates).

• Improved System Capability: The AI exhibits Guaranteed Code and Structure Consistency. When asked to generate complex, multi-step code or coordinate data, it not only produces syntactically correct output but also guarantees that the underlying structure is maximally consistent and predictable, minimizing runtime bugs caused by internal drift.

Improvement 6: Developing a Matched-Scale Distillation Baseline for Low-Resource Scenarios.

• Technical Action: Implement a structured distillation pipeline where the full, large-scale model (the teacher, e.g., Qwen2.5-Coder 32B) is tasked with generating high-quality, verified examples across difficult domains (Iconic/Compositional inversions). These samples are then used to train a smaller, highly efficient student model on a matched scale, ensuring the student retains the full behavioral capability of the teacher without requiring its vast parameter count.

• Rationale: This addresses the necessity of deploying high-capability models in resource-constrained environments while maintaining state-of-the-art performance.

• Improved System Capability: The AI achieves High Performance, Low Latency Deployment. It can execute complex spatial reasoning tasks and generate sophisticated code with minimal computational overhead, making it viable for edge devices or high-throughput API services where latency is critical.

Large language models (LLMs) trained only on text and code can sometimes generate programs that draw recognizable images. However, it is unclear whether this reflects an internal representation of 2D spatial layout or simply the ability to translate spatial descriptions into code. We introduce Autoregressive Mosaics (AM-Bench), a benchmark that separates these factors: First, a translation task gives a model a fully specified geometry of a picture in words as a prompt and asks for the code that produces it. Second, a layout task requires the model to compose an image from an underspecified prompt. Across eight open-weight text-and-code-only models, all models reliably translate specified geometry into code, but their open-ended layout performance differs substantially, indicating that these differences are not explained by code-generation ability alone. An output-medium ablation further shows that the interface or medium of expression that the model uses matters: replacing procedural code with raw SVG improves layout scores across all models. Finally, probing model activations shows that a coarse layout plan is present before generation, but reflects only the layout implied by the prompt. During generation, models track the evolving geometric state instead of executing an initially fixed plan. Overall, these results show that 2D spatial performance in text-only LLMs depends on both the model and the output medium, and is not explained by code-generation ability alone.