"On the Feature Learning in Diffusion Models"

Why diffusion models are better at learning: They don't pick favorites in features

This paper reveals how diffusion models learn features differently from classification models, showing they maintain balanced feature learning while classification models tend to prioritize specific patterns. The research provides theoretical and empirical evidence for this behavior through a novel feature learning framework.

-----

https://arxiv.org/abs/2412.01021

🔍 Original Problem:

→ Despite diffusion models showing exceptional capabilities in various tasks, there's limited understanding of how they learn features compared to traditional classification models.

→ The theoretical foundations of feature learning in diffusion models remain unexplored, particularly in understanding their superior out-of-distribution performance.

-----

🛠️ Solution in this Paper:

→ The paper develops a theoretical framework analyzing feature learning dynamics in both diffusion and classification models.

→ They use a two-layer convolutional neural network with quadratic activation to study signal-to-noise ratio (SNR) effects.

→ The analysis focuses on how models learn balanced representations through denoising objectives versus classification objectives.

→ They validate findings using synthetic and real-world datasets to demonstrate distinct feature learning patterns.

-----

💡 Key Insights:

→ Diffusion models learn features with linear growth initially, maintaining balanced signal-to-noise ratios

→ Classification models show exponential growth in feature learning, focusing on either signal or noise

→ The ratio of signal learning to noise learning in diffusion models is proportional to n·SNR²

→ Feature learning balance explains diffusion models' improved robustness and transferability

-----

📊 Results:

→ Diffusion models achieve stationary points with signal-to-noise learning ratio of Θ(n·SNR²)

→ Classification models show sharp phase transitions based on n·SNR² threshold

→ Empirical validation shows 64% prediction accuracy on test sets

→ Model achieves 2.21 Sharpe ratio on sector rotation strategy