Technical Summary: VisRefiner

1. Problem Statement

The task of Screenshot-to-Code Generation involves translating static user interface (UI) screenshots into executable frontend code (e.g., HTML/CSS) that faithfully reproduces the target layout and style. While recent Multimodal Large Language Models (MLLMs) have demonstrated end-to-end capabilities in this domain, they typically rely on one-way supervised mapping from designs to code. A critical limitation of these existing approaches is that models are trained without observing the visual outcomes of their generated code. Consequently, they lack the ability to learn how specific visual discrepancies relate to necessary code changes.

In contrast, human developers employ an iterative render–compare–revise process: they render an implementation, compare it against the design, identify visual inconsistencies, and refine the underlying code. Current MLLM-based systems do not internalize this mechanism, treating visual differences merely as external diagnostic signals rather than internal learning cues.

2. Methodology: VisRefiner

To address this gap, the authors propose VisRefiner, a training framework designed to enable MLLMs to learn directly from visual differences between rendered predictions and reference designs. The framework shifts the paradigm from feed-forward prediction to difference-driven learning through two distinct stages:

A. Difference-Aligned Supervision (Stage 1)

This stage establishes an explicit correspondence between observed visual discrepancies and the specific code updates required to resolve them.

• Data Construction (VisDiffUI): The authors constructed a dataset, VisDiffUI, containing paired examples of visual deviations and their corresponding code edits. This dataset is derived from public sources (Vision2UI, WebSight) and augmented via two methods:
  1. Rule-based Perturbations: A library of rules modifies HTML/CSS tokens to simulate localized UI inconsistencies across six dimensions: color, layout, alignment, component, image, and text.
  2. Model-Predicted Imperfections: Imperfect outputs from earlier-stage models are combined with ground truths to capture realistic artifacts (e.g., misplaced containers, redundant wrappers) that are difficult to synthesize via rules.
• Training Objective: The model is trained to predict the refined code ($C_{t+1}$) given the target screenshot ($I_{gt}$), the current code ($C_t$), and the rendered image of the current code ($I_t$). The objective minimizes the negative log-likelihood:
  $$L_{SFT} = -\log P_\theta(C_{t+1} | I_t, I_{gt}, C_t)$$
  This grounds the model's understanding of how implementation changes affect visual appearance.

B. Reinforcement Learning with Self-Refinement (Stage 2)

Building on the supervised foundation, this stage enables the model to refine its own outputs through visual feedback using Group Relative Policy Optimization (GRPO).

• Quantifying Improvement: A Relative Improvement Score (RIS) is defined based on CLIP-based perceptual similarity. It measures the relative progress in reducing visual difference between the current rendering and the target, normalized to ensure scale invariance.
• Reward Design: A composite reward function guides the optimization:
  1. Format Reward ($r_{format}$): Penalizes syntactically invalid code.
  2. Refinement Reward ($r_{improve}$): Assigns a positive signal if the new rendering improves visual similarity.
  3. Quality Reward ($r_{quality}$): Uses the continuous RIS value to measure the degree of improvement.
• Optimization: The model generates candidate code variants, renders them, and updates its policy based on the group-relative advantage of the rewards, encouraging iterative reasoning and internalization of refinement behavior.

3. Key Contributions

1. VisRefiner Framework: A novel training paradigm that integrates visual feedback directly into the training process, allowing MLLMs to learn how code modifications influence visual outcomes.
2. VisDiffUI Dataset: A specialized dataset and data pipeline that aligns visual differences with corresponding code edits, providing a resource for studying appearance-to-implementation correspondences.
3. Empirical Validation: Extensive experiments demonstrating that VisRefiner achieves consistent improvements in visual fidelity, layout alignment, and self-improvement capability across multiple benchmarks.

4. Experimental Results

The authors evaluated VisRefiner (specifically a 7B parameter variant) against proprietary models (GPT-4o, Claude 3.5 Sonnet) and open-source baselines (Qwen2.5-VL, MiMo-VL) on Design2Code, Design2Code-HARD, and the held-out VisDiffUI-Test.

• Single-Step Generation: VisRefiner-7B achieved strong single-step accuracy, outperforming the 72B parameter Qwen2.5-VL-72B on the challenging Design2Code-HARD benchmark (85.6 vs. 80.8 average score) and surpassing GPT-4o in structural alignment (Block Match: 78.1 vs. 56.6).
• Self-Refinement Stability: Unlike baselines where refinement often led to inconsistent results (sometimes degrading performance), VisRefiner demonstrated consistent gains across all datasets and refinement modes.
• Ablation Studies:
  • SFT as the Driver: Supervised Fine-Tuning (SFT) on difference-aligned pairs was identified as the primary factor improving structural and visual quality (e.g., Position Alignment and Color Consistency).
  • GRPO as the Refiner: GRPO provided further incremental improvements, particularly in fine-grained metrics, but relied on the difference-aware foundation built by SFT.
• Human Evaluation: Human annotators found that VisRefiner-7B exhibited "deliberate and active refinement," significantly reducing "Tie" cases (where no change was made) compared to base models. While this led to a slightly higher "Lose" rate (indicating active but sometimes unsuccessful edits), it confirmed the model's ability to engage in non-trivial correction rather than conservative inaction.
• Multi-turn Analysis: While performance fluctuated in multi-turn refinement, filtered results (selecting the best output across turns) continued to improve, suggesting the model explores diverse correction paths.

5. Significance and Claims

The paper claims that VisRefiner represents a shift in screenshot-to-code generation from a purely feed-forward prediction task to a difference-driven learning paradigm. By learning from visual differences, the framework endows models with a more human-like reasoning process, where the model understands the causal link between code edits and visual outcomes.

The authors emphasize that this approach does not merely improve generation quality but also instills a self-refinement capability that is often absent in standard MLLMs, especially at smaller scales. The work suggests that incorporating visual discrepancies as direct learning signals, rather than post-hoc feedback, is a critical step toward advancing multimodal reasoning in structured code synthesis.