From RGB Generation to Dense Field Readout: Pixel-Space Dense Prediction with Text-to-Image Models

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-to-token-to-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 p x p x K_t 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.48x faster than an edit-plus-latent-decode counterpart--dense perception can benefit from generative pretraining without inheriting its output interface.

한 줄 요약

ReChannel은 텍스트-이미지 생성 모델을 활용한 밀집 예측에서 VAE 디코더를 제거하고, 토큰-로컬 선형 헤드를 통해 33K 파라미터로 빠르고 정확한 예측을 달성한다.

핵심 기여도

ReChannel: 생성된 토큰 필드의 채널을 RGB 대신 태스크-네이티브 필드로 재해석하는 새로운 밀집 예측 인터페이스 제안.
기존 VAE 디코더 대신 단일 토큰-로컬 선형 헤드(33K 파라미터) 사용, 공간 믹싱 없이도 성능 유지.
FLUX-Klein 기반으로 6개 밀집 예측 태스크에서 최신 기록 달성 (예: KITTI depth, trimap-free matting).
4B 파라미터 규모에서 기존 편집+잠재 디코딩 방식 대비 2.48× 빠르고 정확도 높음.

핵심 아이디어

기존 밀집 예측 접근은 생성 모델의 VAE 디코더를 활용해 태스크 결과를 이미지처럼 디코딩하지만, 이는 밀집 예측에 불필요한 복잡성을 유발한다. ReChannel은 DiT 모델이 이미 이미지 평면에서 토큰-패치-토큰 구조를 통해 RGB를 조직하고 있음을 관찰하고, 이 구조를 그대로 활용해 각 토큰의 채널을 RGB 대신 태스크-네이티브 필드로 재해석한다. 이는 생성 모델의 사전 학습된 공간 구조를 그대로 사용하면서도, 별도의 디코더 없이도 밀집 필드를 추출할 수 있게 한다. 이 접근은 DiT가 이미 RGB 이미지를 패치 단위로 조직하고 있으므로, 각 토큰이 이미 공간적 위치 정보를 포함하고 있음을 활용한 것이다.

기술적 접근법

**모델 구조**: FLUX-Klein 기반 DiT 모델 사용.
**입력 인코더**: VAE 인코더는 유지하여 DiT의 입력 분포를 보존.
**출력 인터페이스**: VAE 디코더는 제거.
**토큰 적응**: LoRA를 사용한 가벼운 파라미터 조정으로 DiT 토큰 필드를 태스크에 맞게 조정.
**출력 맵핑**: 각 토큰을 p × p × K_t 크기의 픽셀-스페이스 패치로 매핑하는 토큰-로컬 선형 헤드 사용.
**파라미터 수**: 약 33K, 공간 믹싱 없이도 성능 유지.

주요 결과

**KITTI depth**: 기존 기록을 넘는 정확도 달성.
**Trimap-free matting**: 최신 기록 수립.
**Referring segmentation**: 경쟁력 있는 성능 유지.
**4B 파라미터 규모에서**: 기존 편집+잠재 디코딩 방식 대비 2.48× 빠르고 정확도 +Z% 개선 (Z는 초록색으로 표기됨).
**다양한 태스크**: 6개 밀집 예측 태스크, 10개 이상의 벤치마크에서 최신 기록 달성.

의의 및 한계

ReChannel은 생성 모델의 사전 학습된 RGB 필드를 밀집 예측에 직접 활용할 수 있음을 보여주며, 기존 디코더 기반 접근의 한계를 극복한다. 특히, VAE 디코더를 제거함으로써 계산 비용을 줄이고, 태스크-네이티브 필드를 직접 추출할 수 있게 함으로써 효율성을 높였다. 그러나 FLUX-Klein 이외의 생성 모델이나, 픽셀-정렬되지 않은 태스크에 대한 확장성은 아직 검증되지 않았다. 또한, 태스크별로 LoRA 조정이 필요하므로, 즉각적인 다중 태스크 학습은 어려울 수 있다.

실용적 활용

ReChannel은 이미지 생성 모델을 기반으로 한 다양한 밀집 예측 태스크(예: 깊이 추정, 세그멘테이션, 매트 추출)에 적용 가능하다. 특히, 생성 모델을 기반으로 한 이미지 편집, AR/VR, 자율 주행, 의료 영상 분석 등에서 빠른 처리와 높은 정확도를 요구하는 상황에 유용하다.