Pre-trained Text-to-Image Diffusion Models Are Versatile Representation Learners for Control

🌈 Abstract

The paper introduces Stable Control Representations (SCR), which are pre-trained vision-language representations extracted from text-to-image diffusion models. The authors show that these representations can outperform state-of-the-art models like CLIP and ViT-based representations on a wide range of embodied control tasks, including manipulation, navigation, and fine-grained visual prediction. The key contributions are:

1. A multi-step approach for extracting versatile vision-language representations from text-to-image diffusion models.
2. Evaluation of diffusion model representations on a broad range of embodied control tasks.
3. Systematic analysis of the design choices for extracting effective representations from diffusion models.

🙋 Q&A

[01] Layer Selection and Aggregation

1. Questions related to the content of the section?

• What design choices did the authors consider for extracting representations from the diffusion model's U-Net?
• How did the authors aggregate the feature maps from different layers of the U-Net?
• Why did the authors choose to concatenate the feature maps from the mid and downsampling blocks of the U-Net?

The authors considered extracting representations from different layers of the diffusion model's U-Net. They found that concatenating the feature maps from the mid and downsampling blocks, and then passing them through a learnable convolutional layer, resulted in a representation size comparable to other pre-trained models while preserving important details. This multi-layer aggregation approach was found to be instrumental to the high performance of the SCR representations.

[02] Diffusion Timestep Selection

1. How did the authors determine the optimal diffusion timestep for extracting representations? The authors hypothesized that control tasks requiring detailed spatial understanding would benefit from using fewer diffusion timesteps, corresponding to a later stage in the denoising process. They experimented with different timestep values and found that performance was sensitive to this choice, with timesteps around 0-10 performing the best on the Franka-Kitchen manipulation tasks.

[03] Prompt Specification

1. How did the authors investigate the role of text prompts in shaping the SCR representations? For tasks with language specifications, the authors provided the corresponding text prompt to the U-Net during representation extraction. For purely vision-based tasks, they explored whether constructing reasonable text prompts would affect downstream policy learning when using the language-guided visual representations from the U-Net.

The authors found that providing text prompts did not consistently improve downstream performance, and in some cases, even degraded performance when the prompts were irrelevant to the visual context.

[04] Intermediate Attention Map Selection

1. How did the authors leverage the cross-attention maps generated by the diffusion model's U-Net? The authors hypothesized that the word-level cross-attention maps generated by the U-Net, which align the visual features with the text embeddings, could help downstream control policies generalize to an open vocabulary of object categories.

They tested this hypothesis on the OVMM open-vocabulary navigation task, where they fused the cross-attention maps with the extracted feature maps from the U-Net (referred to as SCR-attn). This variant outperformed the standard SCR representations, demonstrating the benefits of incorporating the text-aligned attention information.

[05] Fine-Tuning on General Robotics Datasets

1. How did the authors fine-tune the Stable Diffusion model to better align it towards generating representations for control tasks? The authors fine-tuned the base Stable Diffusion model on a small subset of datasets commonly used for representation learning in embodied AI, including EpicKitchens, Something-Something-v2, and Bridge-v2. This fine-tuning was done using the same text-conditioned generation objective as the base model, but without using any task-specific data.

The authors found that this simple fine-tuning approach was effective in most cases, with the fine-tuned SCR-ft representations outperforming the base SCR model on several tasks.

</output_format>