Loss determination for latent diffusion models
Described is a system for improving machine learning models by accessing a first latent diffusion machine learning model, accessing a second latent diffusion machine learning model that was derived from the first latent diffusion machine learning model, the second latent diffusion machine learning model trained to perform a second number of denoising steps, generating noise data, processing the noise data via the first latent diffusion machine learning model to generate one or more first latent features, processing the noise data via the second latent diffusion machine learning model to generate one or more second latent features, and inputting the one or more first latent features and the one or more second latent features into a loss function. The system then modifies a parameter of the second latent diffusion machine learning model based on the output of the loss function.
1 . A system comprising:
at least one processor; and
at least one memory component storing instructions that, when executed by the at least one processor, cause the at least one processor to perform operations comprising:
accessing a first latent diffusion machine learning model, the first latent diffusion machine learning model trained to perform a first number of denoising steps;
accessing a second latent diffusion machine learning model that was derived from the first latent diffusion machine learning model, the second latent diffusion machine learning model trained to perform a second number of denoising steps;
generating noise data;
processing the noise data via the first latent diffusion machine learning model to generate one or more first latent features;
processing the noise data via the second latent diffusion machine learning model to generate one or more second latent features;
inputting the one or more first latent features and the one or more second latent features into a loss function; and
modifying a parameter of the second latent diffusion machine learning model based on an output of the loss function;
wherein the loss function comprises a distillation loss;
wherein the distillation loss comprises:
ℒ
vani
_
dst
1
=
w
(
λ
t
)
x
^
t
(
s
)
-
z
t
″
-
σ
t
″
σ
t
z
t
σ
t
″
-
σ
t
″
σ
t
α
t
2
2
,
wherein
w
(
λ
t
)
=
max
(
α
t
2
σ
t
2
,
1
)
is a truncated SNR weighting coefficient,
wherein w(λ t ) is a truncated signal-to-noise ratio (SNR) weighting coefficient determined as a function of a relative signal-to-noise parameter λ t for a diffusion timestep
t
,
x
t
(
s
)
represents a predicted denoised latent produced by the student latent diffusion machine learning model at the timestep t, Z t″ represents a teacher-model-generated noisy latent at an earlier timestep t″, α t and σ t are timestep-dependent signal and noise scaling coefficients, respectively, determined by a diffusion noise schedule, α t″ and σ t″ are corresponding signal and noise scaling coefficients associated with the timestep t″, and
·
2
2
denotes a squared Euclidean norm computed in latent space.
2 . The system of claim 1 , wherein the loss function further comprises an original loss.
3 . The system of claim 2 , wherein the original loss comprises:
ℒ
ori
=
t
∼
U
[
0
,
1
]
,
x
∼
p
data
(
x
)
,
e
∼
𝒩
(
0
,
1
)
v
^
θ
(
t
,
z
t
,
c
)
-
v
2
2
,
wherein v is a ground-truth target velocity, E t~U[0,1], x~p data (x),ε~N(0,1) denotes an expectation over (i) a diffusion timestep tsampled from a uniform distribution, (ii) a clean latent x sampled from a training data distribution, and (iii) random noise ε sampled from a multivariate normal distribution, θ e (t, z t , c) is a velocity prediction output by the latent diffusion machine learning model parameterized by θ, conditioned on the diffusion timestep t, a noisy latent z t , and an optional conditioning input c, z t is a noisy latent generated by combining the clean latent xwith the random noise εaccording to a diffusion noise schedule, and
·
2
2
computed in latent space.
4 . The system of claim 2 , wherein the loss function includes a random probability distribution, and wherein the loss function applies a first type of distillation loss or a second type of distillation loss based on the random probability distribution.
5 . The system of claim 4 , wherein the first type of distillation loss applies a classifier-free guidance function to determine the distillation loss, wherein the first type is different than the second type.
6 . The system of claim 5 , wherein the second type of distillation loss does not apply a classifier-free guidance function to determine the distillation loss.
7 . The system of claim 1 , wherein modifying the parameter includes performing back propagation on the second latent diffusion machine learning model based on the output of the loss function that processes the one or more first latent features of the first latent diffusion machine learning model and the one or more second latent features of the second latent diffusion machine learning model.
8 . The system of claim 1 , wherein the operations further comprise:
restructuring the first latent diffusion machine learning model to perform a third number of denoising steps, the first number of denoising steps being larger than the third number of denoising steps; and
restructuring the second latent diffusion machine learning model to perform a fourth number of denoising steps, the second number of denoising steps being larger than the fourth number of denoising steps,
wherein processing the noise data via the first latent diffusion machine learning model comprises processing the noise data via the restructured first latent diffusion machine learning model, and
wherein processing the noise data via the second latent diffusion machine learning model comprises processing the noise data via the restructured second latent diffusion machine learning model.
9 . The system of claim 8 , wherein the third number of denoising steps is half the first number of denoising steps.
10 . The system of claim 8 , wherein the fourth number of denoising steps is half the second number of denoising steps.
11 . The system of claim 8 , wherein the fourth number of denoising steps is half the third number of denoising steps.
12 . The system of claim 1 , wherein the first and second latent diffusion machine learning models are stable diffusion models including cross-attention blocks and ResNet blocks.
13 . The system of claim 1 , wherein the second latent diffusion machine learning model is derived from the first latent diffusion machine learning model by restructuring a UNet architecture of the first latent diffusion machine learning model.
14 . The system of claim 13 , wherein restructuring of the UNet architecture includes changing the architecture of cross attention and ResNet blocks.
15 . The system of claim 1 , wherein the first latent diffusion machine learning model includes a denoising architecture, wherein the number of denoising steps include a number of iterations for transmitting an output of a prior iteration as input to a current iteration of the denoising architecture, and wherein processing the noise data via the first latent diffusion machine learning model to generate one or more first latent features includes iteratively processing the noise data via the denoising architecture for the first number of denoising steps to generate first latent features.
16 . The system of claim 1 , the operations further comprising:
receiving a prompt for image generated from a user; and
processing the prompt via the second latent diffusion machine learning model with the modified parameter to generate one or more user-requested images.
17 . The system of claim 16 , wherein the output of the second latent diffusion machine learning model is processed through a decoder to generate an output image of a current iteration.
18 . A method comprising:
accessing, by one or more processors, a first latent diffusion machine learning model, the first latent diffusion machine learning model trained to perform a first number of denoising steps;
accessing a second latent diffusion machine learning model that was derived from the first latent diffusion machine learning model, the second latent diffusion machine learning model trained to perform a second number of denoising steps;
generating noise data;
processing the noise data via the first latent diffusion machine learning model to generate one or more first latent features;
processing the noise data via the second latent diffusion machine learning model to generate one or more second latent features;
inputting the one or more first latent features and the one or more second latent features into a loss function; and
modifying a parameter of the second latent diffusion machine learning model based on an output of the loss function;
wherein the loss function comprises an original loss, wherein the original loss comprises:
ℒ
ori
=
𝔼
t
~
U
[
0
,
1
]
,
x
~
p
data
(
x
)
,
e
~
N
(
0
,
I
)
v
^
θ
(
t
,
z
t
,
c
)
-
v
2
2
,
wherein v is a ground-truth target velocity, E t~U[0,1], x~p data (x),ε~N(0,1) denotes an expectation over (i) a diffusion timestep tsampled from a uniform distribution, (ii) a clean latent x sampled from a training data distribution, and (iii) random noise ε sampled from a multivariate normal distribution, v θ (t, z t , c) is a velocity prediction output by the latent diffusion machine learning model parameterized by θ, conditioned on the diffusion timestep t, a noisy latent z t , and an optional conditioning input c, z t is a noisy latent generated by combining the clean latent xwith the random noise εaccording to a diffusion noise schedule, and
·
2
2
denotes a squared Euclidean norm computed in latent space.
19 . A non-transitory computer-readable storage medium, the computer-readable storage medium including instructions that when executed by a computer, cause the computer to perform operations comprising:
accessing a first latent diffusion machine learning model, the first latent diffusion machine learning model trained to perform a first number of denoising steps;
accessing a second latent diffusion machine learning model that was derived from the first latent diffusion machine learning model, the second latent diffusion machine learning model trained to perform a second number of denoising steps;
generating noise data;
processing the noise data via the first latent diffusion machine learning model to generate one or more first latent features;
processing the noise data via the second latent diffusion machine learning model to generate one or more second latent features;
inputting the one or more first latent features and the one or more second latent features into a loss function; and
modifying a parameter of the second latent diffusion machine learning model based on an output of the loss function;
wherein the loss function comprises a distillation loss and an original loss;
wherein the distillation loss comprises:
ℒ
vani
_
dst
1
=
w
(
λ
t
)
x
^
t
(
s
)
-
z
t
″
-
σ
t
″
σ
t
z
t
σ
t
″
-
σ
t
″
σ
t
α
t
2
2
,
wherein
w
(
λ
t
)
=
max
(
α
t
2
σ
t
2
,
1
)
is a truncated SNR weighting coefficient, wherein w(λ t ) is a truncated signal-to-noise ratio (SNR) weighting coefficient determined as a function of a relative signal-to-noise parameter λ t for a diffusion timestep
t
,
x
t
(
s
)
represents a predicted denoised latent produced by the student latent diffusion machine learning model at the timestep t. Z t″ represents a teacher-model-generated noisy latent at an earlier timestep t″, α t and σ t are timestep-dependent signal and noise scaling coefficients, respectively, determined by a diffusion noise schedule, α t and σ t are corresponding signal and noise scaling coefficients associated with the timestep t″, and
·
2
2
denotes a squared Euclidean norm computed in latent space.