使用Itô密度估計器對擴散模型進行叠加

The Superposition of Diffusion Models Using the Itô Density Estimator

December 23, 2024

作者: Marta Skreta, Lazar Atanackovic, Avishek Joey Bose, Alexander Tong, Kirill Neklyudov

cs.AI

摘要

易於存取的多個預訓練擴散模型的寒武紀大爆發，顯示了對於結合多個不同預訓練擴散模型的方法的需求，而無需承擔重新訓練更大結合模型所帶來的顯著計算負擔。本文將在生成階段將結合多個預訓練擴散模型的問題，置於一個新提出的名為超位置的框架下。從著名的連續方程原理嚴謹地推導出超位置，並設計了兩種專為在SuperDiff中結合擴散模型而量身定制的新算法。SuperDiff利用一種新的可擴展It\^o密度估算器來計算擴散SDE的對數概似，與用於計算分歧的眾所周知的Hutchinson估算器相比，不會產生額外開銷。我們展示了SuperDiff在大型預訓練擴散模型上的可擴展性，因為超位置僅在推論過程中通過組合進行，並且在實現上也非常方便，因為它通過自動重新加權方案結合不同的預訓練向量場。值得注意的是，我們展示了SuperDiff在推論時的效率，並模擬了傳統的組合運算符，如邏輯OR和邏輯AND。我們在實驗中展示了使用SuperDiff在CIFAR-10上生成更多樣化圖像的效用，使用穩定擴散進行更忠實的提示條件圖像編輯，以及改進的蛋白質無條件全新結構設計。https://github.com/necludov/super-diffusion

English

The Cambrian explosion of easily accessible pre-trained diffusion models suggests a demand for methods that combine multiple different pre-trained diffusion models without incurring the significant computational burden of re-training a larger combined model. In this paper, we cast the problem of combining multiple pre-trained diffusion models at the generation stage under a novel proposed framework termed superposition. Theoretically, we derive superposition from rigorous first principles stemming from the celebrated continuity equation and design two novel algorithms tailor-made for combining diffusion models in SuperDiff. SuperDiff leverages a new scalable It\^o density estimator for the log likelihood of the diffusion SDE which incurs no additional overhead compared to the well-known Hutchinson's estimator needed for divergence calculations. We demonstrate that SuperDiff is scalable to large pre-trained diffusion models as superposition is performed solely through composition during inference, and also enjoys painless implementation as it combines different pre-trained vector fields through an automated re-weighting scheme. Notably, we show that SuperDiff is efficient during inference time, and mimics traditional composition operators such as the logical OR and the logical AND. We empirically demonstrate the utility of using SuperDiff for generating more diverse images on CIFAR-10, more faithful prompt conditioned image editing using Stable Diffusion, and improved unconditional de novo structure design of proteins. https://github.com/necludov/super-diffusion