PixelDense: Dense Prediction as Representation Alignment for Pixel Diffusion

REPA matches intermediate denoiser features to a frozen encoder and speeds up diffusion training by an order of magnitude without changing inference. Its targets, however, are almost always semantic encoders such as DINOv2 and CLIP. iREPA showed that the alignment benefit comes mainly from the teacher's spatial structure rather than its global semantics, and dense-prediction foundation models are trained precisely to expose that structure: instance and region boundaries from SAM2, monocular depth from Depth Anything v2, and surface normals from Metric3D v2. We ask whether this semantic and geometric knowledge can supervise an image generator during training, without needing any dense outputs at inference.

Pixel-space diffusion is a natural fit, because the denoiser tokens lie on the same RGB patch grid that dense teachers consume, so each token aligns directly with a spatially matched teacher feature. As Figure 2 shows, the two groups of teachers respond to complementary structure in the same image.

Dense teachers help individually, but compete under a flat sum. On a PixelGen text-to-image backbone, each of SAM2, Depth Anything v2, and Metric3D v2 improves GenEval Overall over the DINOv2-only baseline when added as an alignment target, and the two geometric teachers lead the segmentation teacher. Yet summing all four REPA losses through a single denoiser projection lands below the best single geometric teacher (0.8036 vs. 0.8069). Semantic encoders reward viewpoint invariance and instance identity, while geometric encoders reward metric layout and surface orientation. When both kinds of gradients share one feature basis, the teachers compete instead of compose. Adding more teachers is therefore not enough: teacher composition is a representation-routing problem.

PixelDense reads the token state \(\mathbf{h}_\ell\) at block 8 of a 16-block pixel diffusion denoiser and projects it into two parallel streams. Each stream feeds light per-teacher heads that are aligned to frozen teacher features with a tokenwise cosine loss.

Two streams alone do not guarantee specialization, since both read the same denoiser token. If \(W_\text{sem}\) and \(W_\text{geo}\) learn the same row subspace, the streams compete for the same denoiser coordinates. PixelDense therefore penalizes the mean absolute cosine between every semantic row and every geometric row:

\[ \mathcal{L}_\text{orth}(W_\text{sem}, W_\text{geo}) = \frac{1}{D_s D_g}\,\bigl\|\bar{W}_\text{sem}\,\bar{W}_\text{geo}^{\top}\bigr\|_1, \qquad \bar{W}_{i,:} = \frac{W_{i,:}}{\|W_{i,:}\|_2}. \]

The full objective adds the factored alignment losses and the orthogonality penalty to the backbone's flow-matching loss, with \(\lambda_\text{orth} = 0.01\) and \(\mathcal{L}_\text{aux}\) (LPIPS and Perceptual-DINO) inherited from PixelGen:

\[ \mathcal{L} = \mathcal{L}_\text{FM} + \underbrace{0.5\,\mathcal{L}^\text{F}_\text{DINO} + 0.3\,\mathcal{L}^\text{F}_\text{SAM} + 0.3\,\mathcal{L}^\text{F}_\text{DA2} + 0.3\,\mathcal{L}^\text{F}_\text{M3D}}_{\mathcal{L}_\text{align}} + \lambda_\text{orth}\,\mathcal{L}_\text{orth} + \mathcal{L}_\text{aux}. \]

Only the denoiser, the two projections, and the four teacher heads receive gradients. After training, the teachers, projections, and heads are all discarded, and the deployed model is the original three-channel RGB generator with its original sampler.

We fine-tune the released 1.1B PixelGen-XXL backbone at 512×512 for 10,000 steps (effective batch size 256, 2×H200) on BLIP3-o-60K. Both backbones are pretrained with DINOv2 REPA, which stays active during fine-tuning, so the primary baseline is the matched DINOv2-only fine-tune. PixelDense improves GenEval Overall from 0.7927 to 0.8093, a 2.09% relative gain, along with DPG-Bench and HPS v2.1. The gains concentrate on the compositional axes that touch object boundaries, depth layout, and surface geometry.

Text-to-image metrics do not measure whether a generation preserves the geometry of a specific input scene. We therefore corrupt an input image at noise level \(\tau \in \{0.5, 0.7, 0.9\}\), reconstruct it with its caption, and score the reconstruction with off-the-shelf probes that are independent of the REPA teachers: COCO-trained OneFormer (Swin-L and DiNAT-L) for panoptic PQ and mIoU, and Marigold for depth and surface normals. PixelDense preserves structure better than the matched PixelGen baseline on every dataset, reference type, probe, and noise level. The gap is largest at \(\tau = 0.5\) and narrows at \(\tau = 0.9\), where severe noise removes most of the input structure for both models.

Every added dense teacher improves GenEval Overall over DINOv2 alone, and the geometric teachers lead (Depth Anything v2 0.8069, Metric3D v2 0.8060, SAM2 0.8020). The naive four-teacher sum falls below the best single teacher under all three weightings. Full PixelDense reaches 0.8093 and outperforms all single-teacher and unfactored variants. Either stream alone scores lower, so the semantic and geometric cues are complementary under the factored objective, and the leave-one-out rows show that every teacher contributes. The design ablations favor projection widths that match the native teacher widths, a mid-network alignment block, and a moderate orthogonality weight.

PixelDense is not tied to one backbone. Applying the same teachers, losses, hyperparameters, and checkpoint-selection rule to DeCo at 512×512, without retuning, improves both GenEval Overall and DPG-Bench over the matched DeCo fine-tune. The gain is smaller than on PixelGen, likely because DeCo starts from a stronger model, but the direction is consistent.

We compare PixelGen and PixelDense on matched prompts and seeds, using the same 25-step sampler as the quantitative evaluation. The examples track the quantitative gains: object count, relative position, color binding, object boundaries, and perspective.