Task-oriented Learnable Diffusion Timesteps for Universal Few-shot Learning of Dense Tasks

TL;DR

This work tackles the challenge of heuristic diffusion timestep features in diffusion model based universal few‑shot dense prediction. It introduces Task-aware Timestep Selection (TTS) and Timestep Feature Consolidation (TFC) along with a parameter‑efficient adapter to adapt a pre trained Latent Diffusion Model to unseen tasks, by selecting a set of $k$ timesteps. Evaluations on the Taskonomy Tiny dataset show improvements in semantic segmentation and other dense prediction tasks, validating the universality and efficiency of the approach. The contributions enable robust few‑shot dense prediction with minimal parameter overhead and broad applicability to diffusion model based systems.

Abstract

Denoising diffusion probabilistic models have brought tremendous advances in generative tasks, achieving state-of-the-art performance thus far. Current diffusion model-based applications exploit the power of learned visual representations from multistep forward-backward Markovian processes for single-task prediction tasks by attaching a task-specific decoder. However, the heuristic selection of diffusion timestep features still heavily relies on empirical intuition, often leading to sub-optimal performance biased towards certain tasks. To alleviate this constraint, we investigate the significance of versatile diffusion timestep features by adaptively selecting timesteps best suited for the few-shot dense prediction task, evaluated on an arbitrary unseen task. To this end, we propose two modules: Task-aware Timestep Selection (TTS) to select ideal diffusion timesteps based on timestep-wise losses and similarity scores, and Timestep Feature Consolidation (TFC) to consolidate the selected timestep features to improve the dense predictive performance in a few-shot setting. Accompanied by our parameter-efficient fine-tuning adapter, our framework effectively achieves superiority in dense prediction performance given only a few support queries. We empirically validate our learnable timestep consolidation method on the large-scale challenging Taskonomy dataset for dense prediction, particularly for practical universal and few-shot learning scenarios.

Figures (5)

• Figure 1: Overall training pipeline of our few-shot learner optimized with diffusion timestep features. Our learning framework operates based on the selected timestep features from the denoising diffusion model pre-trained on a large-scale dataset. As opposed to previous works that rely on heuristic and intuitive selection of diffusion timesteps for a specific task, we leverage the visual representation power from the pre-trained latent diffusion model to search, select, and consolidate multiple timestep features to universally predict dense tasks given a few-shot supports.
• Figure 2: Detailed description of our proposed TTS module and TFC module. Our TTS module takes as input $k$ diffusion timestep features and searches for a new timestep at which its corresponding feature contains more semantically meaningful information than the previous. We first identify the least meaningful timestep feature based on the sum of task losses based on Leave-One-Out procedure, and then selects a new timestep feature under the feature similarity constraint and the resulting task loss. We repeat the selection process for $N_{iter}$ steps until we find a new timestep feature. TFC module matches the image timestep features and label features to examine the effectiveness of each timestep, following the general cross attention mechanism.
• Figure 3: Similarity scores for the timestep features from pre-trained SD v1.4.
• Figure 4: Selected timesteps per task. For each task in the ordinate, the selected $k=4$ timesteps are shown in the abscissa.
• Figure 5: Qualitative comparison. We visualize the qualitative results for all ten evaluated tasks in the Taskonomy-Tiny dataset. Ours (last row) mostly exhibits superior prediction results in all the experimented tasks.