Command Palette
Search for a command to run...
From RGB Generation to Dense Field Readout: Pixel-Space Dense Prediction with Text-to-Image Models
From RGB Generation to Dense Field Readout: Pixel-Space Dense Prediction with Text-to-Image Models
Zanyi Wang Xin Lin Haodong Li Dengyang Jiang Yijiang Li
Abstract
Large-scale text-to-image models are attractive backbones for dense prediction because RGB generation pretraining learns rich semantic, structural, and geometric priors. Existing generative and editing approaches reuse these priors by casting dense prediction as target generation: annotations such as depth, normals, alpha mattes, masks, and heatmaps are encoded into an RGB-trained VAE latent space and decoded back as image-like targets. We argue this inherits more of the generative output interface than dense prediction requires: unlike RGB synthesis, dense prediction asks for pixel-correct, task-native fields on the same image plane, not new RGB content to be rendered. Our key observation is that a pretrained DiT already organizes RGB inputs through a patch→token→patch lattice on the image plane, so each token indexes a fixed output patch whose channels can carry task-native quantities instead of RGB appearance. We instantiate this as ReChannel: we keep the VAE encoder for the DiT's input distribution but drop the target-side decoder, adapt the frozen DiT with task LoRA, and map each token to its ?? × ?? × ???? pixel-space patch through a shared token-local linear head—about 33K parameters, no spatial mixing. Using FLUX-Klein, we evaluate on six dense prediction tasks and over a dozen benchmarks. This minimal interface sets new state-of-the-art on trimap-free matting, KITTI depth, and referring segmentation, and stays competitive on normals, saliency, and pose. In a matched 4B setting it is more accurate and 2.48× faster than an edit-plus-latent-decode counterpart—dense perception can benefit from generative pretraining without inheriting its output interface.
One-sentence Summary
ReChannel, proposed by UCSD and HKUST researchers, repurposes a pretrained DiT by replacing its VAE decoder with a lightweight token-to-patch linear head (∼33K parameters) to directly output task-native pixel-space fields, achieving state-of-the-art on trimap-free matting, KITTI depth, and referring segmentation while being 2.48× faster than an edit-plus-latent-decode counterpart in a matched 4B setting.
Key Contributions
- ReChannel reinterprets the output tokens of a frozen pretrained Diffusion Transformer (DiT) as spatial carriers, mapping each token to a fixed pixel patch whose channels carry task-native quantities (e.g., depth, matting alpha) instead of RGB, thereby eliminating the target-side VAE decoder.
- The method uses a minimal, unified interface: a frozen text-to-image DiT backbone, lightweight task-specific LoRA adaptation, and a shared token-local linear head (about 33K parameters) that directly reads out pixel-space patches without any spatial decoder or upsampling, handling geometry, masks, segmentation, and heatmaps in one form.
- Across six dense prediction tasks and more than a dozen benchmarks, ReChannel achieves new state-of-the-art results on trimap-free matting, KITTI depth, and referring segmentation, stays competitive on surface normals, saliency, and pose, and in a matched 4B-parameter setting is more accurate and 2.48× faster than an edit-plus-latent-decode counterpart.
Introduction
RGB-based dense prediction estimates pixel-aligned fields such as depth, masks, mattes, and heatmaps from a single image, a capability central to many vision tasks. Current approaches either rely on discriminative foundation encoders with task-specific decoders or repurpose text-to-image generative models as priors. The generative route has proven effective, but existing methods largely treat dense targets as images to be synthesized: they encode task fields into an RGB-trained VAE latent space and then decode them back to pixels. This target-side rendering interface is indirect, imposing an RGB reconstruction burden on quantities that are not appearance, and it adds computational overhead. The authors argue that the generative prior’s value lies in the rich spatial field it already organizes, not in its ability to generate the target. They introduce ReChannel, a method that drops the target-side VAE. Instead, a frozen diffusion transformer backbone is lightly adapted with task-specific LoRA, and a token-local linear projection directly reads out task-native patch fields from the adapted token grid. This reframes dense prediction as rechanneling the spatial carriers rather than generating an image, yielding a unified, minimal interface that achieves state-of-the-art results across six tasks while being up to 2.48 times faster than an edit-plus-latent-decode counterpart.
Method
The authors propose ReChannel, a framework that repurposes a pretrained text-to-image diffusion transformer for dense prediction tasks without using the generative decoder. The core idea is to treat the spatial token grid of the transformer as a readout substrate, directly projecting each token to a local pixel-space patch via a lightweight linear mapping. This avoids the conventional image-generation interface and keeps the dense targets in their native pixel-aligned format.
The input pipeline retains the original RGB pathway of the pretrained T2I model. An input image x is encoded by the frozen VAE encoder, producing a latent that lies in the pretrained model’s distribution. The latent is then fed into the diffusion transformer (DiT) backbone, which is run in a deterministic zero-noise mode. The backbone weights stay frozen, and task-specific adaptation is injected through low-rank (LoRA) parameters Δt:
Zt=Fθ+Δt(VAEenc(x);σ=0,ct),
where Fθ is the pretrained transformer, ct is the text condition (e.g., a referring expression), and Zt={zijt} is the resulting spatial token field for task t. This design preserves the pretrained RGB field structure while steering it toward target semantics through LoRA.
The key innovation is the readout stage. Instead of a generative decoder or a multi-stage spatial head, each adapted token zijt is mapped independently to a local dense patch using a shared token-local linear projection:
Y^ijt=reshape(Wtzijt+bt),Y^ijt∈Rp×p×Kt.
Here p is the patch size and Kt the number of output channels for the task. The projection is applied identically at every token, with no inter-token spatial mixing. The reshape operation merely unpacks the projection into pixel space; the spatial structure of the prediction is entirely inherited from the adapted token field. This is the ReChannel readout: it reinterprets each token’s channels from RGB appearance to task-native fields while leaving the spatial carrier intact.
The same readout mechanism is instantiated across diverse tasks. For each task, a separate LoRA adapter and readout projection are trained, with Kt and the supervision loss chosen according to the target representation. Scalar fields like depth, alpha mattes, and saliency use Kt=1; surface normals use Kt=3; pose estimation uses multi-channel heatmaps. Continuous regression tasks are trained with standard pixel-space losses, while binary mask tasks use pixel-space mask supervision. The loss functions are kept conventional, as the focus is on the output interface rather than task-specific loss engineering.
Experiment
ReChannel is evaluated on six dense prediction tasks by adapting a frozen FLUX-Klein text-to-image backbone with a LoRA adapter and a token-local linear head, following each benchmark's standard protocol. It achieves state-of-the-art or competitive results across depth, matting, referring segmentation, and pose, while ablations confirm that LoRA adaptation and a strong pretrained prior are necessary, and that the adapted token field already encodes the output, making extra output machinery or target-side VAE decoding redundant. The token-local readout is the most accurate configuration and also the fastest, as it eliminates the decoder entirely.
ReChannel-9B achieves the best depth and surface normal estimates on most benchmarks, surpassing both specialized and generative baselines. For monocular depth, it leads on KITTI and ScanNet while narrowly trailing one method on NYU. On surface normals, ReChannel-9B attains the lowest mean angular error on all three datasets, and its smaller ReChannel-4B variant is already competitive with the strongest generative approaches. ReChannel-9B sets the best absolute relative depth error on KITTI and ScanNet, improving over Edit2Perc on those two datasets. On surface normals, ReChannel-9B obtains the lowest mean angular error on NYU, ScanNet, and iBims, while ReChannel-4B matches the top generative baselines. Generative methods like GeoWizard are consistently outperformed by ReChannel on both depth and normals, with larger gaps on KITTI depth and iBims normals. Specialized discriminative methods (e.g., DAv2, DSINE) are surpassed by ReChannel-9B on all geometric tasks except NYU depth, where DAv2 still holds an edge.
ReChannel-9B sets a new state of the art on trimap-free matting, outperforming the strongest specialized and generative-prior baselines by clear margins. The advantage is even more pronounced in the zero-shot AIM-500 setting, where the error is less than half that of the previous best generative method, indicating improved cross-domain generalization. The token-local readout removes the target-side VAE, which is particularly beneficial for this boundary-sensitive task. ReChannel-9B achieves new state-of-the-art SAD on both P3M-500-P (5.69) and P3M-500-NP (6.67), surpassing the specialized ViTAE-S by 0.55 and 0.92 points respectively, and the generative-prior GenPercept by 4.06 and 6.10 points. On zero-shot AIM-500, ReChannel-9B reaches 34.90 SAD, less than half of GenPercept's 75.5, showing that avoiding a target-side VAE round-trip enhances both in-domain boundary accuracy and out-of-distribution transfer.
ReChannel models achieve state-of-the-art referring segmentation on all RefCOCO splits without any mask-proposal head, LLM decoder, or task-specific architecture. The 4B variant already outperforms all prior 7B–8B LLM-based methods and the previous best PSALM-1.3B, while the 9B model further raises the average cIoU to 82.0. This indicates that a token-local linear readout from an adapted T2I backbone alone provides both the spatial grounding and segmentation precision that earlier work sought through dedicated output machinery. ReChannel-4B surpasses every prior method, including the 1.3B PSALM (77.1 avg cIoU) and 7B–8B LLM-based models, reaching 80.3 avg cIoU. ReChannel-9B is the best on all eight splits and sets a new state of the art at 82.0 avg cIoU, without any dedicated mask head or LLM decoder.
ReChannel-9B reaches 79.2 AP on COCO pose estimation, surpassing the specialized ViTPose-L by 0.9 points without any pose-specific architecture. For saliency detection, it achieves the best Fmax and MAE on DUTS-TE and ECSSD among all reported methods. These results show the same token-local readout adapts to structured heatmaps and binary masks. ReChannel-9B improves COCO AP by 0.9 over ViTPose-L using a generic readout and no pose-specific design. On saliency, ReChannel-9B sets the highest Fmax and lowest MAE on both DUTS-TE and ECSSD.
When the backbone is frozen and only a linear head is trained, surface normals and matting collapse to near-trivial predictions, revealing that the pretrained token field is not natively aligned with task outputs. Using a random initialization instead of the pretrained prior leaves a large accuracy gap, and adding extra output capacity or generative decoding paths does not improve upon a simple token-local linear readout. The token-local readout is the most accurate and fastest configuration, as the target structure is already present in the adapted token field. Head-only training on a frozen body causes surface normal error to exceed 44° and matting SAD to surpass 170, close to a constant-prediction regime. Random initialization of the same architecture yields roughly double the matting error and significantly worse normals, confirming reliance on the strong pretrained prior. A 13× larger convolutional head and full fine-tuning of the 4B backbone do not outperform the token-local linear readout, and generative interfaces (latent, VAE-decoded, edit) are consistently worse. The token-local linear readout adds no measurable latency, while generative alternatives are 1.56× to 2.48× slower due to additional decoding passes.
ReChannel models are evaluated across a wide range of dense prediction tasks, including monocular depth, surface normals, matting, referring segmentation, pose estimation, and saliency, using standard benchmarks and comparisons against specialized and generative baselines. The experiments validate that a simple token-local linear readout from a pretrained diffusion backbone can achieve state-of-the-art results on all these tasks without any task-specific architectures, outperforming both specialized discriminative methods and generative approaches. Ablation studies confirm that the strong pretrained prior is essential and that the token-local readout is both the most accurate and efficient design, as the target spatial structure is already captured in the adapted token field.