Universal Few-shot Spatial Control for Diffusion Models

We gratefully appreciate your review, acknowledging “the matching module is a good method to unify different representations of label conditions”. Here, we address the raised concerns and questions.

Weaknesses

W1. Motivation for Layout Control-Specific Design

A1. We appreciate the reviewer’s comment. While spatial control is a subset of image-to-image generation, few-shot adaptation for spatial control poses unique challenges, requiring specialized designs. Many image-to-image generation tasks, such as image restoration, image editing, and style transfer, involve a (near) natural image as a condition. Since distribution shifts in the condition are generally not significant in these tasks, it has been observed that a large-scale pre-training is often sufficient to achieve reasonable generalization performance across tasks. However, the spatial control involves spatial guidance signals as conditions, such as edge/depth maps, which are generally non-natural images and have unique distributions across tasks. As shown in our experiments (Table 1 and 2) and Table G in our response to reviewer G2yU, this causes existing image-to-image generation models to fail to generalize in unseen spatial conditions in most cases. UFC is designed specifically to handle the large discrepancy between different spatial condition types, such as encoding (unseen) control conditions by the interpolation of image features through matching, parameter-efficient adaptation mechanisms equipped with meta-training, each of which turned out to be crucial in few-shot adaptation of unseen control tasks (Table 3 and our response to weakness 1 raised by reviewer d7Cj). We will clarify these motivations in the paper.

W2 Differences in the matching module between UFC and Chameleon

A2. While we agree with the reviewer that UFC and Chameleon share a high‑level design—both use a matching module to draw task‑relevant information from a support set for few‑shot adaptation—their objectives and guiding intuitions are fundamentally different.

UFC employs patch‑wise matching to unify diverse condition types by extracting task‑agnostic features from support images. Chameleon [1], in contrast, uses matching to harvest task‑specific features from support labels (which correspond to our support conditions).

Because the goals diverge, the roles of the attention components are effectively reversed. In UFC, condition features from a task‑specific encoder serve as the Query and Key, weighting task-agnostic image features from a shared encoder that play the Value role; this yields unified control features. In Chameleon, image features from a task‑specific encoder act as the Query and Key to express task‑specific similarity, while label features from a shared encoder act as the Value, producing a task‑specific map that a shared label decoder then processes.

Finally, UFC conditions its matching module on the diffusion timestep embedding, allowing control features to adapt at every denoising step, whereas Chameleon does not include timestep conditioning.

W3 Task-specific parameters selection; clarification on condition encoders; computing resources.

A3.

(1) (Task-specific parameters $\theta_\tau$ selection): While any parameter-efficient fine-tuning method can be applied to implement the task-specific parameters $\theta_\tau$, we adopt bias-tuning (as described in L188-189) as it has been shown to be effective in few-shot dense prediction literature [1,2]. In this implementation, the bias parameters to be tuned are the portion of the parameters in the encoder.

(2) (clarification on condition encoders): We use a single, shared condition encoder for encoding the conditions (both the support and query), as stated in Eq.(3). We apologize for the misleading figure and will modify the condition encoder part to make it clearer.

(3) (computing resources): We report the computing resources and hyperparameters for test-time finetuning in Table A and B, respectively.

Table A. Computing Resources

Table B. Hyperparameters

W4 Reason why fine-tuning with Prompt-Diffusion overfits on certain conditions

A4. To clarify, PromptDiffusion also slightly overfits to Depth. Therefore, PromptDiffusion overfits on the high-level conditions, such as depth, normal, pose, and densepose, while not on the low-level conditions, such as Canny and HED. We conjecture that this is due to our 3-fold evaluation protocol, where we split the tasks into three folds and use the other folds for meta-training when evaluating each fold. Since low-level labels can be relatively easily inferred from high-level labels, while the opposite is hard, we suspect that PromptDiffusion shows better generalization on low-level tasks by transferring its knowledge from high-level tasks contained in the meta-training dataset.

On the other hand, this behavior is not observed in UFC, which highlights the effectiveness of the robust patch-wise matching and parameter-efficient fine-tuning for few-shot adaptation to unseen tasks.

W5 About the ablation study on the matching module

A5. We want to clarify that our ablation study tests the matching module’s role in terms of the adaptive knowledge extraction from the support sets, instead of the matching layer itself. The simple linear projection is not capable of adaptively extracting relevant information for each query condition from the support set, and cannot cope with different support sizes (e.g, increased shot after fine-tuning). Furthermore, the addition of all features obtained after the linear projection can destroy the structural information encoded in the feature of the query condition and potentially amplify over-fitting.

To validate our claim, we conducted an additional experiment following the suggestion of the reviewer. Table C and D summarize the results. As expected, we observe that the w/o matching, w/ support variant (V1) fails to yield meaningful controllability and also damages the image quality a lot. This indicates that the ablation variant in our paper (V2) serves as a more reasonable baseline to ablate the effect of matching.

Table C. Controllability Measurement on COCO 2017 Validation

Table D. FID (↓) Evaluation on COCO 2017 Validation

W6. More explanation on the relevance between query/reference patches in Supp. Fig.5

A6. We would like to clarify that query patches attend to relevant support patches, not necessarily semantically identical ones. Since there is no explicit training objective for the matching module to model the semantic similarity, it is not very surprising that “knee” (we believe that the reviewer may miswrite it as elbow) patch can be attended to patches of other joints. We conjecture that what the condition encoder learned is to produce condition features that capture general “human joints”. Then the matching on these features composes a skeleton-like layout of features, which seems to be sufficient to guide the diffusion model. Similarly, for Canny/Depth, patches are matched by spatial layout rather than semantics. We hope our discussion on attention maps resolves the concern of the reviewer.

W7. Performance saturation analysis with increased sample size

A7. We provide performances of UFC with more samples (up to 300) in Table E and Table F. The results show the same trend as our previous discussion in Section 5.3 of the paper (L296 to L299): with more support data, the model gains better controllability performance. The controllability seems to be saturated for most of the tasks, except for Pose. We hypothesize that Pose requires semantic understanding of human joints given sparse labels, thus more challenging to learn than other tasks, and requires much more data to reach the saturation point.

Table E. Controllability Evaluation on COCO 2017 Validation Set

Table F. FID (↓) Evaluation on COCO 2017 Validation Set

[1] Kim, Donggyun, et al. "Chameleon: A data-efficient generalist for dense visual prediction in the wild." ECCV 2024.

[2] Kim, Donggyun, et al. "Universal Few-shot Learning of Dense Prediction Tasks with Visual Token Matching." ICLR 2023.

Thank you for the response. The response is thorough and in detail. After reading all reviews, I still think the link between the motivation and the method is very weak. In the response, the authors explain the difference of general image-to-image synthesis and layout-control generation. I want to know which part of the proposed method addresses the specific challenge of this layout-control generation task, to distinguish it from the general I2I generation.

We sincerely appreciate your review and recognition of our “novel framework” that addresses “high data requirements and limited adaptability” of current spatial-conditioning diffusion models. We address the raised concern below.

Weaknesses

W1. Additional explanations on support set and computing resources

A1. The support set is a set containing several condition-image pairs. For each pair, we refer to its condition and image as support condition and support image, respectively. During meta-training, we sample the support set randomly for each task at every iteration to simulate different few-shot episodes, while the support set is given and fixed during the fine-tuning and inference (Section 4.3). To save computation, we employ three-shot supports in each training episode and five-shot supports during inference. The meta-training takes 12 hours using 8 RTX 3090 GPUs, while fine-tuning takes less than an hour using a single RTX 3090 GPU. When conditioned on a five-shot support set, inference takes about 3 seconds to process an image using a single RTX 3090 GPU.

W2. Comparison with more recent controllable generation methods

A2. To our best knowledge, there have been recent advances in the image-to-image generation tasks, but it was hard to find comparable baselines. These methods fall into several categories, such as architecture improvement [1-5] for spatial control, or unifying architecture for image-to-image generation tasks [6 - 11]. However, none of these claim to have strong few-shot adaptation as well as zero-shot transfer for spatial control. For instance, OmniGen [6] demonstrates few-shot learning on unseen classes for segmentation, but also acknowledges that it fails to handle unseen input types (normal prediction). As a demonstration, we also conducted experiments to evaluate the zero-shot and one-shot capability of OmniGen and the zero-shot capability of Flux.1 Kontext [10] (please refer to Table G in our response to reviewer G2yU) and find that this model fails or shows limited performance in spatial control on unseen tasks.

In that regard, we would appreciate it if the reviewer could suggest any baseline so that we can compare.

W3. Additional performance verification on real-world few-shot scenarios

A3. We appreciate the reviewer for the suggestion. We would like to kindly remind the reviewer that we have the results for the requested experiments in Figure 10 and Section F of the supplementary file. In this experiment, we employed three unseen spatial control tasks from FreeControl, such as point clouds, meshes, and wireframes, which are inherently data scarce and difficult to obtain model-based annotations, as suggested by the reviewer. Similar to FreeControl, we were able to report only the qualitative results due to the absence of pre-trained models to measure quantitative controllability.

As shown in Figure 10, our method can produce convincing results with only a few-shot support (30 examples), although these unseen spatial conditions are significantly different from the ones in the meta-training dataset (Canny, HED, Depth, Normal, Pose, Densepose). We appreciate the reviewer’s suggestion and will include this result in the main paper.

[1] Tan, Zhenxiong, et al. "Ominicontrol: Minimal and universal control for diffusion transformer." ICCV 2025.

[2] Tan, Zhenxiong, et al. "Ominicontrol2: Efficient conditioning for diffusion transformers." arXiv preprint arXiv:2503.08280 (2025).

[3] Yu, Fanghua, et al. "UniCon: Unidirectional Information Flow for Effective Control of Large-Scale Diffusion Models." ICLR 2025.

[4] Zavadski, Denis, et al. "Controlnet-xs: Rethinking the control of text-to-image diffusion models as feedback-control systems." ECCV, 2024.

[5] Li, Ming, et al. "Controlnet++: Improving conditional controls with efficient consistency feedback” ECCV 2024.

[6] Xiao, Shitao, et al. "Omnigen: Unified image generation." CVPR 2025.

[7] Han, Zhen, et al. "Ace: All-round creator and editor following instructions via diffusion transformer." arXiv preprint arXiv:2410.00086 (2024).

[8] Lin, Yijing, et al. Realgeneral: Unifying visual generation via temporal in-context learning with video models. ICCV 2025.

[9] Lin, Weifeng, et al. "Pixwizard: Versatile image-to-image visual assistant with open-language instructions." ICLR 2025.

[10] Black Forest Lab. "FLUX. 1 Kontext: Flow Matching for In-Context Image Generation and Editing in Latent Space." arXiv preprint arXiv:2506.15742 (2025).

[11] Wang, Haoxuan, et al. "Unicombine: Unified multi-conditional combination with diffusion transformer." arXiv preprint arXiv:2503.09277 (2025).

We greatly appreciate your review, recognizing our proposed framework addressing “an important challenge in controllable generation”. Here, we address your concerns and questions.

Weaknesses

W1. Performance of UFC without the meta-training stage

A1. We would like to clarify that, in the few‑shot learning literature [1–3], it is common to call a setting K‑shot when only K image–label pairs are available for each unseen task, even though auxiliary data may be used for pre‑training or meta‑training. Our references to 30‑shot follow this convention. Because our goal is few‑shot generalisation to unseen condition types, meta‑training is needed to give the universal control adapter prior knowledge of diverse conditions.

To verify this, we skipped meta‑training and adapted UFC directly to the unseen conditions, still using 30 support examples. Results for three representative tasks are shown below. As expected, omitting meta‑training degrades performance, confirming that this stage is crucial for universal few‑shot control.

Table A. Controllability measurement on COCO 2017 validation

Table B. FID (↓) evaluation on COCO 2017 validation

W2. Data fairness of the comparison across baselines

A2. We would like to clarify that all few-shot baselines in Table 1 and 2 (Prompt Diffusion, Uni-ControlNet+FT, Prompt Diffusion+FT) are pre-trained using the exact same meta-training dataset of our method and fine-tuned on the same few-shot data, as briefly discussed in Ln 221 - 225 in our paper. For fully-supervised methods, ControlNet was trained for each condition separately following previous practice [4, 5, 6] since it is designed to work on a single condition, hence cannot be pre-trained from other tasks. Considering that ControlNet still leverages full supervision of 150k examples for each task, but UFC leverages only 30 examples out of it to learn new spatial control tasks, we believe ControlNet serves as a strong upper-bound for various tasks as reported in the previous works. We also note that we considered the multi-task variants of ControlNet, Uni-ControlNet, which leverages full supervision of 6 evaluation tasks jointly for training. Compared to Uni-ControlNet, UFC was meta-trained with 4 tasks of Uni-ControlNet training data, and fine-tuned with only 30 examples for each new task, irrelevant to meta-training tasks. Therefore, UFC has no advantage over both few-shot and fully-supervised baselines in terms of data access. We appreciate the comment and will make this description clearer in the paper.

Questions

Q1. About diversity and selection strategy for support set

A1.

During the few‑shot fine‑tuning stage reported in our main paper, we randomly select 30 support examples for Canny, HED, Depth, and Normal. For Pose and DensePose, however, we manually curate 30 examples to ensure sufficient diversity for these higher‑level semantic tasks.

Because our matching strategy hinges on condition-specific similarity, the notion of “diversity” depends on tasks. For edge inputs such as Canny and HED, the support set must span a wide range of edge-map densities and orientations. For Pose and DensePose, it needs to cover varied human scales, occlusions, and pose deformations. In that regard, for most of the tasks in our experiments (except Pose and Densepose), the object class (e.g., cats or dogs) is irrelevant to ensure the diversity of the support set. For instance, our model could generalize to COCO 2017 validation set with only 30 support examples, which can cover only a small set of objects.

To measure how support choice influences performance, we repeated fine-tuning three times, each time drawing a new random set of 30 supports and carrying out few-shot generation. The results are summarized in the tables below.

For all tasks except Pose and DensePose, performance shows no significant sensitivity to support randomness. Pose and DensePose, however, performed worse than the results in Tables 1 and 2 of the main paper, which used curated supports maximising diversity in scale, deformation, crowding, and occlusion. These results confirm that the importance of support-set diversity is itself task-dependent, and having a data selection strategy to ensure sample diversity can potentially benefit the performance (as shown in Pose and Densepose).

Table C. Controllability evaluation with different support sets on COCO 2017

Table D. FID (↓) evaluation with different support sets on COCO 2017

[1] Vinyals, Oriol, et al. "Matching networks for one shot learning." NIPS 2016.

[2] Finn, Chelsea, et al. "Model-agnostic meta-learning for fast adaptation of deep networks." International conference on machine learning. ICML 2017.

[3] Snell, Jake, et al. "Prototypical networks for few-shot learning." NIPS 2017.

[4] Zhao, Shihao, et al. "Uni-controlnet: All-in-one control to text-to-image diffusion models." NeurIPS 2023.

[5] Wang, Zhendong, et al. "In-context learning unlocked for diffusion models." NeurIPS 2023.

[6] Qin, Can, et al. "Unicontrol: A unified diffusion model for controllable visual generation in the wild." NeurIPS 2023.

We sincerely appreciate the review and your acknowledgement of our “clear”, “intuitive” architecture and our “good results”. We address the raised concerns and questions below.

Weaknesses

W1. Clarification on meta-learning and few-shot finetuning setting

A1. We thank the recommendation. We will clarify it in the abstract and introduction.

W4. User studies

We thank the recommendation. Due to the time limit, we will conduct user studies and incorporate the results in our paper after the rebuttal period ends.

Questions

Q1. Clarification on the matching module

A1. There is a separate matching module for each layer of the U-Net encoder, which operates on the features at each level, and its output is incorporated into the corresponding U-Net decoder layer. We will clarify this in the method section.

W2, Q2. Performance with fewer shots; fine-tuning iterations and overfitting; Effect of support set’s diversity on model performance

A2.

(1) (Performance with fewer shots): Following the reviewer’s suggestion, we performed a fewer‑shot ablation; results are in the tables below. Consistent with Figure 4, controllability and quality drop as the supports shrink. However, it also shows that UFC can achieve a certain level of controllability using unseen control conditions, even with a very limited support (5 and 15 shots), showcasing the generalizability of our method.

Table A. Controllability evaluation on COCO 2017

Table B. FID (↓) evaluation on COCO 2017

(2) (Fine-tuning iterations and overfitting): With 30 supports, < 600 steps (< 1hour on an RTX‑3090) suffice for all six tasks. Continuing the fine-tuning over a much longer iterations eventually leads to overfitting, but it was easily avoidable by applying early stopping using a single held-out validation example to monitor the denoising loss. We found that this strategy leads to similar early-stopping iterations across the tasks, with trends that low-level conditions (e.g., edges) tend to converge faster than high-level conditions (e.g., pose).

(3) (Effect of support sets’ diversity on model performance): Because our matching strategy hinges on condition-specific similarity, the notion of “diversity” depends on tasks. For edge inputs such as Canny and HED, the support set must span a wide range of edge-map densities and orientations. For Pose and DensePose, it needs to cover varied human scales, occlusions, and pose deformations. In that regard, for most of the tasks in our experiments (except Pose and Densepose), the object class (e.g., cats or dogs) is irrelevant to ensure the diversity of the support set. For instance, our model could generalize to COCO 2017 validation set with only 30 support examples, which can cover only a small set of objects.

To measure how support choice influences performance, we repeated fine-tuning three times, each time drawing a new random set of 30 supports and carrying out few-shot generation. The results are summarized in the tables below.

Table C. Controllability evaluation with different support sets on COCO 2017

Table D. FID (↓) evaluation with different support sets on COCO 2017

For all tasks except Pose and DensePose, performance shows little sensitivity to support randomness. Pose and DensePose, however, the model performed worse than the results in Tables 1 and 2 of the main paper, which used curated supports maximising diversity in scale, deformation, crowding, and occlusion. These results confirm that the importance of support-set diversity is itself task-dependent.

W3, Q3. Generalization on difficult-to-obtain/in-the-wild conditions

A3. We would like to kindly remind the reviewer that we have the results for the requested experiments in Figure 10 and Section F of the Appendix. In this experiment, we employed three unseen spatial control tasks from FreeControl, such as point clouds, meshes, and wireframes, which are inherently data scarce and difficult to obtain model-based annotations, as suggested by the reviewer. Similar to FreeControl, we were able to report only the qualitative results due to the absence of pre-trained models to measure controllability.

As shown in Figure 10, our method can produce convincing results with only a few-shot support (30 examples), although these unseen spatial conditions are significantly different from the ones in the meta-training dataset. We appreciate the reviewer’s suggestion and will include this result in the main paper.

Q4. Comparison with training-free, fully-supervised baselines

A4. We would like to kindly remind the reviewer that Appendix E already explains why FreeControl is excluded: running it on the full COCO-2017 validation set (5,000 images) is prohibitively slow. The original FreeControl paper reports results on only 30 images from 10 classes—insufficient for fair comparison. We appreciate the suggestion to add Ctrl-X as a fast, training-free baseline. Accordingly, we evaluated Ctrl-X on COCO-2017 using the same Stable Diffusion 1.5 backbone as UFC. As shown in Tables F and G, Ctrl-X consistently underperforms UFC in both controllability and FID across all conditions.

Table E. Controllability measurement on COCO 2017 validation

Table F. FID (↓) evaluation on COCO 2017 validation

We agree that T2I-Adapter would be a valuable fully supervised reference. However, due to our resource limitations, we selected the ControlNet and Uni-ControlNet as the main baselines. Because our goal is to demonstrate few-shot capability rather than to surpass fully supervised methods, we believe the current baseline set is sufficient to support our claims.

Q5. Usefulness and applications for UFC within language-driven/interactive image generation models.

A5. Open-source, language-driven image-generation systems such as Flux.1 Kontext and OmniGen can handle multiple conditioning signals, but only those seen during training. OmniGen’s authors themselves note that “previously unseen image types (e.g., surface-normal maps) can hardly be processed as expected,” and Flux.1 Kontext likewise makes no claim of zero-shot adaptability—the model is trained on natural images rather than abstract control inputs. UFC, in contrast, is designed to adapt to new condition types from just a handful of image–condition pairs. To demonstrate this, we evaluate each model on 30 images from the ImageNet-R-TI2I dataset [1]. Since OmniGen was trained on Canny, HED, Depth, Pose, and Segmentation, we evaluate it on the unseen Normal conditions in both zero-shot and one-shot. Similarly, we zero-shot evaluate Flux.1 Kontext on unseen HED and Normal conditions.

Table G. Controllability evaluation on ImageNet-R-TI2I

Despite OmniGen’s larger and more diverse training corpus, it lags behind UFC, and Flux.1 Kontext performs poorly on both tasks—often simply reconstructing the input condition when we investigate the generated images. These results underline the limited few-shot capacity of current unified generators and the advantage of UFC’s universal few-shot control adapter, which is independent of the diffusion backbone. Because UFC is orthogonal to recent unified image-to-image frameworks, it can be integrated with them to extend their versatility to truly unseen control domains.

[1] Tumanyan, Narek, et al. "Plug-and-play diffusion features for text-driven image-to-image translation." CVPR 2023.

Limitation

L1 Social concern on the training dataset

A1. We acknowledge the societal concern about using LAION‑400M and appreciate the recommendation of the cleaned RE‑LAION‑5B dataset. We were not aware of these issues during the submission period while seeking an open-source dataset. We will add this discussion and consider re‑training UFC on safer data before releasing code and checkpoints.