School of Engineering and Applied Sciences, Harvard University, Cambridge, MA

# Learning Predictive and Interpretable Timeseries Summaries from ICU Data

Nari Johnson  MSc  Sonali Parbhoo  PhD  Andrew S Ross  PhD  Finale Doshi-Velez
PhD

Abstract

Machine learning models that utilize patient data across time (rather than
just the most recent measurements) have increased performance for many risk
stratification tasks in the intensive care unit. However, many of these models
and their learned representations are complex and therefore difficult for
clinicians to interpret, creating challenges for validation. Our work proposes
a new procedure to learn summaries of clinical timeseries that are both
predictive and easily understood by humans. Specifically, our summaries
consist of simple and intuitive functions of clinical data (e.g. “falling mean
arterial pressure”). Our learned summaries outperform traditional
interpretable model classes and achieve performance comparable to state-of-
the-art deep learning models on an in-hospital mortality classification task.

##  1 Introduction

Accurate predictions of patient risk in critical care units can aid clinicians
in making more effective decisions. Specifically, early identification of
patients at high risk for in-hospital mortality is critical to assess patient
disease acuity and inform life-saving interventions [1, 2]. To predict in-
hospital mortality risk, researchers have developed algorithms ranging from
simple score-cards [1, 3] to statistical machine learning (ML) models. Recent
advances in ML have led to the development of models with vast improvements in
predictive accuracy for patient in-hospital mortality risk [4, 5, 6, 7].

Despite these improvements, however, ML models are still prone to critical
errors, often failing to generalize across different care settings or
institutions [8]. An emerging line of research in interpretability, defined by
[9] as “the ability to explain or to present in understandable terms to a
human,” provides an alternative way to ensure systems preserve properties such
as safety or nondiscrimination. If a model is interpretable to stakeholders,
then clinical experts can inspect the model and verify that its reasoning is
sound. This ability to audit and validate is especially important when the
models are used to inform critical decisions affecting patient health.

In this work, we present a novel ML method to learn clinical timeseries
summaries that are interpretable and predictive. We introduce functions to
compute summaries that align with simple and intuitive concepts, such as
whether the timeseries is decreasing or spikes above a critical threshold. In
contrast to prior work, our method learns how much of the timeseries should be
used to calculate these summaries, discarding earlier timesteps that may be
irrelevant for a specific prediction task. Importantly, we introduce
relaxations of our summary definitions to enable differentiable optimization,
allowing summary parameters to be learned jointly with those of a downstream
model. We show that with our method, we can achieve accuracies comparable with
state-of-the-art baselines without sacrificing interpretability.

##  2 Related Work

Prior work on explaining clinical timeseries models fall into two categories:
learning simple models that are inherently interpretable, and generating
explanations of complex black box models. We summarize a few key examples
below.

Interpreting Deep Models. One popular strategy for explaining ML models is
learning a second post-hoc explanation model to explain the first black box
model [10]. Many post-hoc explanation techniques for clinical timeseries
models train explanation models that quantify the relative importance of each
clinical variable [11, 12]. However, several works argue against the use of
post-hoc explanation techniques, as explanation models are not always faithful
or representative of the true underlying black boxes [13, 14]. Our method
avoids these problems by design, instead explicitly optimizing for
interpretability so that a second explanation model is not needed.

Another line of research proposes attention mechanism models specifically
designed for timeseries. [15, 16] present neural attention architectures for
clinical timeseries and argue that attention scores measure feature
importance. However, attention methods are often highly complex and nonlinear.
Furthermore, [17] shows attention scores do not always reflect true
importance. Instead of approximating importance, our study uses linear models
over richer features, where importance does not need to be approximated but
can be directly read off model coefficients.

Expert Systems and Expert Features. Our work extends a long tradition of
clinical experts hand-crafting features to create interpretable clinical
decision-support algorithms. Two expert systems widely used in ICUs are SAPS-
II [1] and APACHE [3], which use simple score-card algorithms to evaluate
patient acuity. These systems use input features such as the patient’s average
or worst lab or vital values over time to compute mortality risk. While SAPS-
II and APACHE are simple and simulable to clinicians, their predictiveness is
limited, as they cannot capture how labs or vitals change across a patient’s
stay.

A similar line of research proposes manual construction of expert features
from clinical data, which are then used as input to ML models [18, 19]. One
limitation of this approach is that expert feature derivation is expensive and
requires clinical expertise. Rather than relying on expert knowledge to
identify which summary features will be the most predictive for a given task,
our work instead uses optimization with a sparsity constraint to automatically
learn which summary features are the most predictive.

Summarizing Clinical Timeseries. A number of works have proposed a wide range
of summary statistics for patient timeseries data. Many works such as [4, 20]
train clinical models using the minimum value, maximum value, first measured
value, or skew of clinical timeseries data. [21] proposes a more comprehensive
set of 14 summary statistics to characterize the central tendency, dispersion
tendency, and distribution shape of clinical timeseries data. Our work extends
this research, and is the first to our knowledge to use the slope of the
timeseries or proportion of time above or below critical thresholds. Our work
is also novel in that we explicitly model and optimize for the duration over
which we compute each summary feature.

##  3 Cohort and Problem Set-Up

Our goal is to learn summaries of patient timeseries data that are both human-
interpretable and predictive for a downstream classification task. Our
approach will define how to calculate these human-interpretable summaries, and
describe how both summary and classification model parameters are learned
through optimization. In what follows, we detail each of these processes.

Prediction Task. Our work examines early prediction of in-hospital mortality.
We use the patient’s first 24 hours of data to predict if they would later
expire over the course of the remainder of their admission. Patients who
expired in the first 24 hours of their stay were excluded from our cohort.

Cohort Selection. We use data from the MIMIC-III critical care database [22],
which contains deidentified health data from patients in critical care units
of the Beth Israel Deaconess Medical Center between 2001 and 2012. All data
was extracted from MIMIC-III PhysioNet version 1.4, which contains 30,232
patients. We exclude patients under 18 years of age and patients whose weight
is not measured. We include data from each patient’s first hospitalization
only, and only patients with stays between 24-72 hours [4]. After applying
these criteria, our final cohort contained 11,035 patients, 15.23% of whom
died in-hospital. Cohort characteristics and demographics are summarized in
Table 1.

Cohort | Age | % Female | % Urgent | % Emergency | % Elective | % MICU | % SICU | % CCU | % CSRU |
---|---|---|---|---|---|---|---|---|---|---
All | 64.7 | 43.8 | 1.12 | 84.46 | 14.43 | 42 | 18 | 12 | 16 |
+ | 70.9 | 46.5 | 0.77 | 96.25 | 2.97 | 53 | 18 | 13 | 4 |
- | 64.1 | 43.6 | 1.14 | 83.22 | 15.64 | 41 | 18 | 12 | 17 |

Table 1: Mean statistics for the population cohort, and for cohorts of
positive versus negative patients for in-hospital mortality. Abbreviations:
MICU, medical care unit; SICU, surgical care unit; CCU, cardiac care unit;
CSRU, cardiac-surgery recovery unit.

###  3.1 Extracting Inputs and Outputs

For each patient $n$ in our $N$ patient cohort, we extracted static
observations and physiological data including labs and vital signs sampled
hourly. All clinical variables were separately normalized to have zero mean
and unit variance. Figure 1 shows how features are extracted for patients that
are positive versus negative for in-hospital mortality.

Figure 1: Example positive and negative time-series to illustrate feature
extraction. The two trajectories have input data extracted from time 0 to time
of prediction $T=24$. Figure inspired by Sherman et al [23].

Static observations $\bm{S}$. Matrix $\bm{S}:(N\times 8)$ contains $8$
demographic variables for each patient $n$: their age at admission, gender,
and other information about their ICU stay (their first ICU service type, and
whether their admission was urgent, emergency, or elective).

Per-timestep clinical observations $\bm{X}$. The clinical variable tensor
$\bm{X}:(N\times D\times T)$ contains $D=28$ measurements of clinical
variables for each of the $N$ patients at time $t$, discretized by hour. These
28 measurements consist of vital signs and labs: diastolic blood pressure,
fio2, GCS score, heartrate, mean arterial blood pressure, systolic blood
pressure, SRR, oxygen saturation, body temperature, urine output, blood urea
nitrogen, magnesium, platelets, sodium, ALT and AST, hematocrit, po2, white
blood cell count, bicarbonate, creatinine, lactate, pco2, glucose, INR,
hemaglobin, and bilirubin. Missing values at timestep $t$ were imputed using
either the most recent measurement of the variable, or the population median
if the variable had not yet been measured during the patient’s stay. We use
the patient’s first $T=24$ hours of data to predict in-hospital mortality. We
use subscripts to index into the tensor: for example, $\bm{X}_{t}$ indicates
the $(N\times D)$ matrix of measurements taken at time $t$.

Per-timestep measurement indicators $\bm{M}$. The measurement indicator tensor
$\bm{M}:(N\times D\times T)$ contains indicator elements $\bm{M}_{n,d,t}$
which are 1 if their corresponding clinical variable in $\bm{X}_{n,d,t}$ was
measured at time $t$, 0 otherwise.

Outcome labels $\bm{y}$. Label vector $\bm{y}$ contains indicators
$\bm{y}_{n}$ which are 1 if patient $n$ expired in-hospital, 0 otherwise.

##  4 Methods

Given a cohort of $N$ training examples $\\{\bm{X},\bm{M},\bm{S},\bm{y}\\}$,
we propose a novel procedure for learning predictive human-interpretable
timeseries summaries. Concretely, we first compute summaries $\bm{H}$ from
clinical timeseries data $(\bm{X},\bm{M})$. We then use the summaries $\bm{H}$
in addition to static data $\bm{S}$ and clinical variables $\bm{X}$ as input
to a Logistic Regression model to predict labels $\bm{y}$. This process of
using human-interpretable summaries for prediction is shown in Figure 2.

In Section 4.1, we motivate and introduce our novel summary features. In
Section 4.2, we give continuous relaxations of summary feature definitions to
enable efficient inference of summary parameters. In Section 4.3, we discuss
how predictive summary features can be jointly learned with downstream
classification model parameters.

Figure 2: Illustrates summary extraction for prediction from timeseries data.
First, non-static clinical variables $\\{\bm{X},\bm{M}\\}$ are used to compute
interpretable summary features $\bm{H}$. Then, summary features $\bm{H}$,
static features $\bm{S}$, and the non-static variables $\\{\bm{X},\bm{M}\\}$
at the time of prediction are given as input to a Logistic Regression model
$g$, which predicts output labels $\bm{\hat{y}}$. Figure inspired by Ghassemi,
Szolovits et al [24].

###  4.1 Model: Defining human-interpretable summaries

Measures of central tendency and dispersion have commonly been used to
summarize timeseries (see survey in [21]). Our key modelling innovations
include adding additional features that correspond to how clinicians
themselves describe timeseries. Our novel contributions include explicitly
modelling the overall trend of a lab/vital and the number of hours that a
lab/vital dips above or below a threshold, as well as allowing different
features to be computed over different periods of time (e.g. the most recent 6
hours vs. the most recent 24 hours).

The $I=13$ summary statistics used in this study are listed in Table 2. Each
of the summary statistics takes into account measurement indicators $\bm{M}$
so that clinical variable summaries are computed only using timesteps where
the variable is measured. Each summary statistic is applied to each of the $D$
clinical variables to create the summary feature tensor $\bm{H}:(N\times
D\times I)$.

Below, we expand on the parameterization of our summary features. Next, we
describe how we enable efficient, automated search over summary feature
parameters. Importantly, our approach automates many processes associated with
summary design, enabling optimization over summary parameters.

Incorporating Duration. Many prior works in clinical timeseries modelling do
not use all timesteps for the patient, but instead only the most recent
available data, such as the six or twelve hours before the time of prediction.
In contrast to prior works, we explicitly model how much of each timeseries
should be used to compute each of the $I$ summaries in $\bm{H}$. For example,
we may wish to exclude earlier measurements of a particular clinical variable
if only the variable’s recent measurements before time of prediction $T$ are
significant for a prediction task. Specifically, for each variable $d$ and for
each summary function $i$, we define a duration time $\bm{C}_{d,i}$. Only the
variable’s timeseries data that occurred in the immediately previous
$\bm{C}_{d,i}$ hours before the time of prediction is used to calculate
summary $i$. We organize all the duration time parameters $\bm{C}_{d,i}$ into
a $(D\times I)$ matrix $\bm{C}$.

To exclude data that occurs before time $(T-\bm{C}_{d,i})$ when computing
summary features, we multiply each of the original timeseries variables by
indicator variables for whether the measurements occurred within
$\bm{C}_{d,i}$ hours before time of prediction $T$. For example, a mean
summary statistic would be computed using indicator variables
$\mathbbm{1}(\cdot)$ as:

$\bm{H}_{i=mean}:(N\times
D)=\left(\textstyle\sum_{t=1}^{T}\mathbbm{1}(t>T-\bm{C}_{i=mean})\odot\bm{X}_{t}\odot\bm{M}_{t}\right)/\left(\textstyle\sum_{t=1}^{T}\bm{M}_{t}\odot\mathbbm{1}(t>T-\bm{C}_{i=mean})\right)$
(1)

where $\odot$ is the element-wise multiplication operator and division is
performed element-wise. Our objective is to learn duration parameters $\bm{C}$
that maximize the predictiveness of their corresponding summary features
$\bm{H}$.

Threshold Parameters. Some of the summary functions $f_{i}$ in Table 2 have
additional parameters such as thresholds. For example, one of the summaries is
the proportion of the patient’s measured timeseries where their measured
clinical variables are above some $D$-dimensional critical threshold parameter
vector $\bm{\phi}^{+}$ for each variable:

$\left(\textstyle\sum_{t=1}^{T}\bm{M}_{t}\odot\mathbbm{1}(\bm{X}_{t}>\bm{\phi}^{+})\right)/\left(\textstyle\sum_{t=1}^{T}\bm{M}_{t}\right)$
(2)

These summaries correspond to clinically-intuitive ideas, such as whether the
patient been mostly well or sick. As with durations, we learn threshold
parameters automatically to avoid burdening experts and to assist in
prediction.

###  4.2 Continuous Relaxations for Efficient Inference

Learning Duration Parameters. In Equation 1, we showed how summary functions
$f_{i}$ that only use the most recent $\bm{C}$ hours of data can be calculated
using indicator variables $\mathbbm{1}(t>T-\bm{C}_{d,i})$. These indicator
variables, however, do not have informative gradients and are not
differentiable. To enable differentiable optimization for duration time
parameters $\bm{C}$, we introduce weight parameters $\bm{W}$ by relaxing the
indicator random variables using the sigmoid function
$\sigma(x)=\frac{1}{1+e^{-x}}$. Using the duration parameter matrix $\bm{C}$,
we define $D$-dimensional vectors $\bm{w}_{t,i}$ that compose weight tensor
$\bm{W}:(T\times I\times D)$:

$\bm{w}_{t,i}=\sigma\left((t-T+\bm{C_{i}})/\tau\right)$ (3)

For each feature $d$ and summary $i$, clinical observations $\bm{X}_{d,t}$
where $t>T-\bm{C}_{d,i}$ will have corresponding weights $\bm{w}_{t,i,d}$ near
1. Timesteps $t$ where $t<T-\bm{C}_{d,i}$ will have corresponding weights
$\bm{w}_{t,i,d}$ near 0. Temperature parameter $\tau$ controls the harshness
of the weight matrix: small temperatures push the sigmoid function towards its
edges, learning weights that are closer to exactly 0 for timesteps before
$T-\bm{C}_{d,i}$ and 1 for timesteps after $T-\bm{C}_{d,i}$, effectively
functioning as the indicator variables in Equation 1.

Weighted summary functions $f_{i}$ used to derive human-interpretable
summaries $\bm{H}$ can be found in Table 2. Duration parameters $\bm{C}$ that
determine weight tensor $\bm{W}$ are included in $\beta_{\bm{H}}$, the set of
all parameters necessary to compute the summaries.

Description | Function
---|---
Mean of the time-series | $\left(\sum_{t=1}^{T}\left(\bm{w}_{t,i}\odot\bm{X}_{t}\odot\bm{M}_{t}\right)\right)/\left(\sum_{t=1}^{T}\bm{M}_{t}\odot\bm{w}_{t,i}\right)$
Variance of the time-series | $\frac{(\sum_{t=1}^{T}\bm{M}_{t}\odot\bm{w}_{t,i})^{2}}{(\sum_{t=1}^{T}\bm{M}_{t}\odot\bm{w}_{t,i})^{2}-\sum_{t=1}^{T}\bm{M}_{t}\odot\bm{W}^{2}_{t}}\odot\sum_{t=1}^{T}\bm{M}_{t}\odot\bm{w}_{t,i}\odot\left(\bm{X}_{t}{-}\bar{\bm{X}}\right)^{2}$
Indicator if feature was ever measured | $\sigma\left(\left(\sum_{t=1}^{T}\bm{w}_{t,i}\odot\bm{M}_{t}\right)/\left(\tau\odot\sum_{t=1}^{T}\bm{w}_{t,i}\right)\right)$
Mean of the indicator sequence | $\left(\sum_{t=1}^{T}\bm{w}_{t,i}\odot\bm{M}_{t}\right)/\left(\sum_{t=1}^{T}\bm{w}_{t,i}\right)$
Variance of the indicator sequence | $\left(\frac{(\sum_{t=1}^{T}\bm{w}_{t,i})^{2}}{(\sum_{t=1}^{T}w_{t})^{2}-\sum_{t=1}^{T}\bm{W}^{2}_{t}}\right)\sum_{t=1}^{T}\bm{w}_{t,i}\odot(\bm{M}_{t}-\bar{\bm{M}})^{2}$
# switches from missing to measured | $\left(\sum_{t=1}^{T-1}\bm{w}_{t,i}\odot|\bm{M}_{t+1}-\bm{M}_{t}|\right)/\left(\sum_{t=1}^{T}\bm{w}_{t,i}\right)$
First time the feature is measured |  $\min t$ s.t. $\bm{M}_{t}=1$
Last time the feature is measured |  $\max t$ s.t. $\bm{M}_{t}=1$
Proportion of time above threshold $\bm{\phi}^{+}$ | $\left(\sum_{t=1}^{T}\bm{w}_{t,i}\odot\bm{M}_{t}\odot\sigma(\frac{\bm{X}_{t}-\bm{\phi}^{+}}{\tau})\right)/\left(\sum_{t=1}^{T}\bm{M}_{t}\odot\bm{w}_{t,i}\right)$
Proportion of time below threshold $\phi^{-}$ | $\left(\sum_{t=1}^{T}\bm{w}_{t,i}\odot\bm{M}_{t}\odot\sigma(\frac{\bm{\phi}^{-}-\bm{X}_{t}}{\tau})\right)/\left(\sum_{t=1}^{T}\bm{M}_{t}\odot\bm{w}_{t,i}\right)$
Slope of a L2 line |  $\frac{\sum_{t=1}^{T}\bm{w}_{t,i}(t-\bar{t}_{\bm{w}})(\bm{X}_{t}-\bar{\bm{X}}_{\bm{w}})}{\sum_{t=1}^{T}\bm{w}_{t,i}(t-\bar{t}_{\bm{w}})^{2}}$, where $\bar{t}_{\bm{w}}=\frac{\sum_{t}\bm{w}_{t,i}\cdot t}{\sum_{t}\bm{w}_{t,i}}$ and $\bar{\bm{X}}_{\bm{w}}=\frac{\sum_{t}\bm{w}_{t,i}\odot\bm{X}_{t}}{\sum_{t}\bm{w}_{t,i}}$
Standard error of the L2 line slope | $1/\left(\sum_{t}\bm{w}_{t,i}\odot(t-\bar{t}_{\bm{w}})^{2}\right)$
Table 2: Table of functions $f_{i}$ used to calculate human-interpretable
summaries $\bm{H}$. For each of the $D$ clinical variables, all $I$ of the
above functions are applied to each of the $N$ patients. Each of the $I$
summary features $i$ is defined with respect to $D$-dimensional weight vectors
$\bm{w}_{t,i}$ defined in Section 4.2. Paramater $\tau$ is a temperature
parameter for the sigmoid function. $\mathbbm{1}(\cdot)$ denotes indicator
variables for events inside the parentheses, $\odot$ indicates element-wise
matrix multiplication and division is done element-wise. Additionally,
$\bar{\bm{X}}=\sum_{t}^{T}\bm{M}_{t}\odot\bm{X}_{t}/\sum_{t}^{T}\bm{M}_{t}$,
and $\bar{\bm{M}}_{t}=\frac{1}{T}\sum_{t}^{T}\bm{M}_{t}$.

Learning Threshold Parameters. Our work relaxes summary definitions to enable
differentiable optimization to learn summary parameters. The indicator
variables used to define the proportion of hours that a patient’s timeseries
is above thresholds $\bm{\phi}^{+}$ in Equation 2 are non-differentiable with
respect to $\bm{\phi^{+}}$. To enable differentiable optimization, our work
defines our threshold summary features using the sigmoid function $\sigma$
with temperature parameter $\tau$:

$f_{threshold}(\bm{X},\bm{M},\bm{W})=\left(\textstyle\sum_{t=1}^{T}\bm{w}_{t,i=threshold}\odot\bm{M}_{t}\odot\sigma\left(\frac{\bm{X}_{t}-\bm{\phi}^{+}}{\tau}\right)\right)/\left(\textstyle\sum_{t=1}^{T}\bm{M}_{t}\cdot\bm{w}_{t,i=threshold}\right)$
(4)

Threshold parameters $\\{\bm{\phi}^{+},\bm{\phi}^{-}\\}$ are included in
$\beta_{\bm{H}}$, the set of all parameters necessary to compute the
summaries.

###  4.3 Learning Process

Our study uses summary features $\bm{H}$, along with static variables $\bm{S}$
and timeseries variables $\\{\bm{X},\bm{M}\\}$ at prediction time $T=24$ as
input to a Logistic Regression model $g$ with coefficients $\bm{\beta}_{g}$.
Logistic Regression model $g$ outputs predicted probabilities
$\hat{\bm{y}}=p(\bm{y}=\bm{1}|\bm{X},\bm{S},\bm{M})$. Our objective is to
learn optimal summary and model parameters
$\bm{\beta}=\\{\bm{\beta}_{\bm{H}},\bm{\beta}_{g}\\}$. We jointly learn
parameters $\bm{\beta}$ by minimizing the loss function

$\mathcal{L}(\bm{\beta};\bm{X},\bm{M},\bm{S},\bm{y})=-\frac{1}{N}\sum_{n=1}^{N}\omega_{n}\left(\bm{y}_{n}\cdot\log[g(\bm{X}_{n},\bm{M}_{n},\bm{S}_{n},\bm{\beta})]+(1-\bm{y}_{n})\log[1-g(\bm{X}_{n},\bm{M}_{n},\bm{S}_{n},\bm{\beta})]\right)+\Omega(\bm{\beta}_{g})$
(5)

Our loss function is the sum of the weighted binary cross-entropy loss using
predictive model $g$ and the regularization penalty $\Omega(\bm{\beta}_{g})$.
To account for class imbalance, we reweight each training example $n$’s loss
contribution by $\omega_{n}$, the inverse of its class frequency in the
training dataset.

Horseshoe Regularization on Coefficients. We explicitly optimize for sparsity
in model parameters $\bm{\beta}_{g}$ via our regularization penalty in our
training objective. We use a Horseshoe regularization penalty
$\Omega(\bm{\beta}_{g})$ with shrinkage parameter $1$ to encourage sparsity in
the learned regression coefficients[25].

##  5 Experiments

We compare models learned using our summaries to other interpretable models as
well as deep learning baselines. We show that our models outperform other
traditional human-interpretable model classes and achieve performance
comparable to deep models on the in-hospital mortality task.

Model configurations. Our baseline models are: Ridge Logistic Regression
models that take as input only the patient timeseries measured at the time of
prediction $T$, and Ridge Logistic Regression and LSTM models that take as
input all of the patient timeseries. We used an LSTM as our deep baseline
architecture as prior works document their superior performance at mortality
prediction from clinical timeseries data [26].

Our LSTM models were trained on the sequential timeseries data
$\\{\bm{X},\bm{M}\\}$ with a step size of 1 hour. The LSTM hidden states at
each timestep were then used to predict both the next timestep of the patient
timeseries $\bm{X}$, and to predict outcome labels $\bm{y}$. All $T$ of the
output hidden states $\bm{X}_{t}$ were input to a fully-connected layer, which
output predictions of the next timestep $\bm{X}_{t+1}$. The last output hidden
state at time $T$ was also input with static features $\bm{S}$ to a fully-
connected layer to predict outcome labels $\bm{y}$. The LSTM models were
trained to minimize both the mean-squared error of the next state prediction,
and the binary cross entropy loss of the classification prediction. ReLU
activation functions were applied to both fully-connected layers.

Training Details. For training and testing, we split the cohort of $N$
patients into train and test sets, where all data associated with each patient
is either in train or test. All performance metrics are averaged across five
train-test splits. Ridge baseline models were implemented using RidgeCV from
scikit-learn [27]. LSTMs as well as our summary-based Logistic Regression
models were implemented with PyTorch, and trained with the Adam optimizer[28]
at a batch size of 256. We trained all of our Logistic Regression models for
30,000 epochs and LSTM models for 10,000 epochs using early stopping. All
hyperparameters (including the LSTM hidden state and layer dimensions,
optimizer learning rate, and regularization parameters) were selected via
random hyperparameter search [29]. All temperature parameters $\tau$ were set
to $0.1$.

Our final learned LSTM models have hidden state dimension 32, 1024 nodes in
the layer to predict the next timestep, and 64 nodes in the layer to predict
labels $\bm{y}$. They are trained using a learning rate of $1e-05$. Our final
learned models with summaries use $\alpha=1e-05$, Horseshoe shrinkage
parameter $1.0$, and learning rate $1e-05$.

##  6 Results

Our learned models achieve performance comparable to state-of-the-art
baselines. Table 3 compares the performance of our learned models with linear
and deep baselines for the in-hospital mortality prediction task. Applying a
linear model to the learned summary features $\bm{H}$ consistently improves
AUCs in comparison to using a linear model on clinical timeseries and static
data alone. Our models achieve an AUC performance comparable to state-of-the-
art LSTM models with test AUCs of $0.9000\pm 0.0223$. Notably, our models that
allow differentiable optimization to learn duration times outperform models
that compute all summary features using the entire duration of the patient
timeseries. This implies that there is predictive value in explicitly
modelling how much of each variable’s clinical timeseries should be considered
for a specific prediction task.

Model | Train set AUC | Test set AUC
---|---|---
LR trained on $\\{\bm{S},\bm{M}\\}$ and $\bm{X}$ at time $T$ only | $0.8653\pm 0.0013$ | $0.8626\pm 0.0079$
LR trained on $\\{\bm{S},\bm{M},\bm{X}\\}$ | $0.8931\pm 0.0015$ | $0.8668\pm 0.0122$
LSTM trained on $\\{\bm{S},\bm{M},\bm{X}\\}$ | $\bf{0.9101\pm 0.0018}$ | $\bf{0.9000\pm 0.0223}$
Our model, trained on $\\{\bm{S},\bm{M},\bm{X}\\}$, non-differentiable durations $\bm{C}$ | $\bf{0.9065\pm 0.0019}$ | $\bf{0.8818\pm 0.0063}$
Our model, trained on $\\{\bm{S},\bm{M},\bm{X}\\}$, differentiable durations $\bm{C}$ | $\bf{0.9074\pm 0.0016}$ | $\bf{0.8867\pm 0.0061}$
Table 3: Performance of learned models on the in-hospital mortality prediction
task. AUCs are averaged over five train-test splits with their standard error.
Abbreviations: LR, Logistic Regression; LSTM, long short-term memory. $\bm{C}$
refers to the duration parameters defined in Section 4.2.

Our learned models use fewer features to achieve higher accuracy in comparison
to other interpretable baselines. To evaluate the sparsity of each model, we
performed a set of ablation experiments where we zeroed all but the $N$
coefficients with the largest magnitudes for each of the learned Logistic
Regression models. In Figure 3, we show the average test set AUC for Logistic
Regression baseline versus our models using only $N$ coefficients. Our models
consistently outperform baseline models when all but $N$ coefficients are
zeroed, suggesting that our models learn a smaller and more predictive set of
important features.

Figure 3: Prediction quality (mean test AUC) vs. model complexity (number of
non-zero features) for baseline Ridge Regression trained on all timesteps of
the patient timeseries, versus our models. The dark blue horizontal line shows
our model’s average test set AUC using all $401$ derived features, and the
orange horizontal line shows the average test set AUC of the Logistic
Regression model trained only on the $65$ features extracted at the time of
prediction.

Our learned models are interpretable. Table 4 shows 15 key summary features
that consistently have the largest learned Logistic Regression coefficients
across train-test splits. The corresponding coefficients for each feature can
be interpreted as a measure of the feature’s contribution to the final
classification label. For example, because the mean of the patient’s GCS has a
large negative coefficient, this means that patients with higher mean GCS
scores will be assigned lower predicted probabilities for in-hospital
mortality. Therefore our models are _decomposable_ [30], as each of the
model’s features and coefficients has an intuitive clinical explanation.

Our learned summary features are clinically sensible. The vast majority of the
key summary features learned by our models shown in Table 4 are supported by
studies in medical literature. For instance, it is widely accepted that
patients who are older tend to have lower chances of survival in ICU settings
[31, 32]. Similarly, patients with lower GCS scores of below 6 tend to have
severe injuries and higher chances of mortality [33]. Notably a lower GCS
score in the later hours of a patient’s hospitalisation significantly reduces
a patient’s chances of survival [34, 35]. Finally, the normal range of
features such as the blood oxygen saturation (SPO2) is between 95% and 100%.
An SPO2 consistently below 90% indicates hypoxaemia or potential respiratory
distress. These patients have to be mechanically ventilated in ICU and
frequently have lower chances of survival [36, 37].

Feature | Aggregation | Time | Coefficient
---|---|---|---
Age | static value | - | 14.9
BUN | value at | hour 24 | 6.2
GCS | mean over | hours 5 - 24 | \- 5.59
HR | value at | hour 24 | 4.69
FiO2 | value at | hour 24 | 4.31
Hct | times measured over | hours 2 - 24 | \- 3.81
HR, | mean over | hours 2 - 24 | 3.78
GCS | value at | hour 24 | \- 3.62
GCS | hours below 6.08 | hours 7 - 24 | 3.37
Creatinine | hours below 0.35 mg/dL | hours 5 - 24 | 2.76
FiO2 | hours above 62.96% | hours 16 - 24 | 2.59
SpontaneousRR | mean over | hours 2 - 24 | 2.29
GCS | times measured over | hours 5 - 24 | 2.23
Sodium | hours below 131.57 mEq/L | hours 1 - 24 | 2.16
WBC | hours below 0.78 cells/mL | hours 6 - 24 | 2.10
SPO2 | hours below 92.36% | hours 10 - 24 | 2.04
Table 4:  Key summary features, sorted from largest to smallest coefficient
magnitudes, from learned models.

Initialization sensitivity. In general, we observed that our optimization
procedure is stable, learning the same 15 key summary features across
different stochastic parameter initializations and train-test splits. However,
there are cases where we observed that the learned duration parameters
$\bm{C}$ varied depending on their initialization. As such, we recommend that
practitioners incorporate prior knowledge about the clinical prediction task
when initializing the duration time parameters. For example, if examining the
entire duration of the patient’s timeseries is necessary for a prediction
task, then the duration parameters should be initialized to include the entire
timeseries by default.

##  7 Discussion & Conclusion

In this work, we defined functions to compute interpretable, parameterizable
summaries of clinical timeseries, and developed relaxations so that our
summary parameters could be jointly learned with a downstream predictive
model. In our experiments, we used Logistic Regression to make predictions
because its coefficients are easily decomposable [30]. However, because our
learned summaries are inherently interpretable, any other interpretable
architecture could be used instead. Our methodology is generalizable and
enables the efficient learning of intuitive and predictive timeseries
summaries without placing any assumptions on the downstream model
architecture.

Future work. Our study poses many interesting directions for future work. One
avenue would be to conduct a user study to validate the human-interpretability
and decomposability of our proposed summary features. Another would be to
evaluate whether the summary features learned for particular critical care
prediction tasks remain predictive for a wider set of critical care prediction
tasks. Finally, we could also develop additional summary statistic functions,
or expand our framework to consider sharing duration parameters across
features or across summaries to better model dependencies between clinical
labs and vitals—as many physiological events are characterized by several
simultaneous changes to multiple labs and vitals [38].

Conclusion. In this paper, we propose a new method to learn interpretable and
predictive summary features from clinical timeseries data. In addition to
introducing novel summary statistics including slope and threshold features,
our work differs from prior work by learning the duration of timeseries data
that should be used to compute each summary. We demonstrate that our learned
timeseries summaries achieve performance quality comparable to state-of-the-
art deep models when trained to predict early patient mortality risk on real
patient data. We also qualitatively validate our models to confirm their
interpretability and sensibility. Our work is an important step towards
optimizing for representations of clinical timeseries data that are both
highly predictive and interpretable.

Acknowledgements: NJ and FDV acknowledge support from NIH R01 MH123804-01A1.
SP acknowledges support from the Miami Foundation and SNSF P2BSP2-184359.

## References

  * [1] JR Le Gall, S Lemeshow, and F Saulnier.  A new simplified acute physiology score (saps ii) based on a european/north american multicenter study.  JAMA, 1993.
  * [2] G Escobar, V Liu, A Schuler, B Lawson, J Greene, and P Kipnis.  Automated identification of adults at risk for in-hospital clinical deterioration.  N Engl J Med, 2020.
  * [3] WA Knaus, DP Wagner, EA Draper, JE Zimmerman, M Bergner, and PG Bastos et al.  The apache iii prognostic system. risk prediction of hospital mortality for critically ill hospitalized adults.  Chest, 1991.
  * [4] A Awad, M Bader-El-Den, J McNicholas, and J Briggs.  Early hospital mortality prediction of intensive care unit patients using an ensemble learning approach.  IJMI, 2017.
  * [5] M Ghassemi et al.  A multivariate timeseries modeling approach to severity of illness assessment and forecasting in icu with sparse, heterogeneous clinical data.  AAAI, 2015.
  * [6] A Rajkomar et al.  Scalable and accurate deep learning for electronic health records.  Digital Medicine, 2018.
  * [7] S Purushothama et al.  Benchmarking deep learning models on large healthcare datasets.  J.Biomed.Inform., 2018.
  * [8] J Futoma et al.  The myth of generalisability in clinical research and machine learning in health care.  Lancet DH, 2020.
  * [9] Doshi-Velez and Kim.  Towards a rigorous science of interpretable machine learning.  arXiv:1702.08608, 2017.
  * [10] MT Ribeiro, S Singh, and C Guestrin.  “why should i trust you?” explaining the predictions of any classifier.  In ACM SIGKDD, 2016.
  * [11] S Tonekaboni et al.  What went wrong and when? instance-wise feature importance for time-series black-box models.  In NeurIPS, 2020.
  * [12] S Lundberg et al.  Explainable machine-learning predictions for the prevention of hypoxaemia during surgery.  Nature Biomedical Engineering, 2018.
  * [13] C Rudin.  Stop explaining black box machine learning models for high stakes decisions and use interpretable models instead.  arXiv:1811.10154v3, 2019.
  * [14] D Shack, S Hilgard, E Jia, S Singh, and H Lakkaraju.  Fooling lime and shap: Adversarial attacks on post hoc explanation methods.  AIES, 2020.
  * [15] E Choi, M Taha Bahadori, J Kulas, A Schuetz, W Stewart, and J Sun.  Retain: An interpretable predictive model for healthcare using reverse time attention mechanism.  arXiv:1608.05745v4, 2017.
  * [16] Y Sha and M Wang.  Interpretable predictions of clinical outcomes with an attention-based recurrent neural network.  ACM BCB, 2017.
  * [17] S Serrano and N Smith.  Is attention interpretable?  arXiv:1906.03731v1, 2019.
  * [18] J Sun, J Hu, D Luo, M Markatou, F Wang, and S Edabollahi.  Combining knowledge and data driven insights for identifying risk factors using electronic health records.  AMIA, 2012.
  * [19] K Roe et al.  Feature engineering with clinical expert knowledge: A case study assessment of machine learning model complexity and performance.  PLoS ONE, 2020.
  * [20] H Harutyunyan et al.  Multitask learning and benchmarking with clinical time series data.  arXiv:1703.07771v3, 2019.
  * [21] C Guo, M Lu, and J Chen.  An evaluation of time series summary statistics as features for clinical prediction tasks.  BMC Medical Informatics Decis. Mak., 20(1):48, 2020.
  * [22] AEW Johnson and TJ Pollard et al.  Mimic-iii, a freely accessible critical care database.  Scientific Data, 2016.
  * [23] E Sherman et al.  Leveraging clinical time-series data for prediction: A cautionary tale.  AMIA, 2018.
  * [24] M Ghassemi, M Wu, MC Hughes, P Szolovits, and F Doshi-Velez.  Predicting intervention onset in the icu with switching state space models.  AMIA, 2017.
  * [25] Hrayr Harutyunyan.  The horseshoe-like regularization for feature subset selection.  Nature, 2019.
  * [26] H Harutyunyan et al.  Multitask learning and benchmarking with clinical time series data.  Scientific Data, 6(1):96, 2019.
  * [27] F Pedregosa et al.  Scikit-learn: Machine learning in Python.  JMLR, 2011.
  * [28] D Kingma and J Ba.  Adam: A method for stochastic optimization.  arXiv:1412.6980, 2014.
  * [29] J Bergstra and Y Bengio.  Random search for hyper-parameter optimization.  JMLR, 2012.
  * [30] Z Lipton.  The mythos of model interpretability.  arXiv:1606.03490v3, 2017.
  * [31] AB Nielsen et al.  Survival prediction in intensive-care units based on aggregation of long-term disease history and acute physiology: a retrospective study of the danish national patient registry and electronic patient records.  The Lancet Digital Health, 2019.
  * [32] L Fuchs et al.  Icu admission characteristics and mortality rates among elderly and very elderly patients.  Intensive care medicine, 38(10):1654–1661, 2012.
  * [33] PG Bastos et al.  Glasgow coma scale score in the evaluation of outcome in the intensive care unit: findings from the acute physiology and chronic health evaluation iii study.  Critical care medicine, 21(10):1459–1465, 1993.
  * [34] RL Sacco et al.  Nontraumatic Coma: Glasgow Coma Score and Coma Etiology as Predictors of 2-Week Outcome.  Archives of Neurology, 1990.
  * [35] C Settervall et al.  In-hospital mortality and the glasgow coma scale in the first 72 hours after traumatic brain injury.  Revista latino-americana de enfermagem, 2011.
  * [36] M Lazzerini et al.  Hypoxaemia as a mortality risk factor in acute lower respiratory infections in children in low and middle-income countries: systematic review and meta-analysis.  PLoS One, 10(9):e0136166, 2015.
  * [37] ML Vold, U Aasebø, T Wilsgaard, and H Melbye.  Low oxygen saturation and mortality in an adult cohort: the tromsø study.  BMC pulmonary medicine, 15(1):1–12, 2015.
  * [38] R Hotchkiss et al.  Sepsis and septic shock.  NRDP, 2016, 2016.


(eccv) Package eccv Warning: Package ‘hyperref’ is loaded with option
‘pagebackref’, which is *not* recommended for camera-ready version

11institutetext: Yandex 22institutetext: Skolkovo Institute of Science and
Technology 33institutetext: Moscow State University 44institutetext: Higher
School of Economics

# YaART: Yet Another ART Rendering Technology

Sergey Kastryulin 1122 Artem Konev∗ 11 Alexander Shishenya∗ 11
Eugene Lyapustin∗ 11 Artem Khurshudov∗ 11 Alexander Tselousov 11
Nikita Vinokurov 11 Denis Kuznedelev 1122 Alexander Markovich 11
Grigoriy Livshits 11 Alexey Kirillov 1133 Anastasiia Tabisheva 11
Liubov Chubarova 1144 Marina Kaminskaia 11 Alexander Ustyuzhanin 11
Artemii Shvetsov 11 Daniil Shlenskii 1122 Valerii Startsev 11
Dmitrii Kornilov 11 Mikhail Romanov 11
Artem Babenko† 11 Sergei Ovcharenko† 11 Valentin Khrulkov† 11

###### Abstract

In the rapidly progressing field of generative models, the development of
efficient and high-fidelity text-to-image diffusion systems represents a
significant frontier. This study introduces YaART, a novel production-grade
text-to-image cascaded diffusion model aligned to human preferences using
Reinforcement Learning from Human Feedback (RLHF). During the development of
YaART, we especially focus on the choices of the model and training dataset
sizes, the aspects that were not systematically investigated for text-to-image
cascaded diffusion models before. In particular, we comprehensively analyze
how these choices affect both the efficiency of the training process and the
quality of the generated images, which are highly important in practice.
Furthermore, we demonstrate that models trained on smaller datasets of higher-
quality images can successfully compete with those trained on larger datasets,
establishing a more efficient scenario of diffusion models training. From the
quality perspective, YaART is consistently preferred by users over many
existing state-of-the-art models.

###### Keywords:

Diffusion models Scaling Efficiency

**footnotetext: Equal contribution$\dagger$$\dagger$footnotetext: Equal senior
authors

##  1 Introduction

Figure 1: RL-aligned YaART generates visually pleasing and highly consistent
images.

Large-scale diffusion models have recently achieved unprecedented success in
text-conditional image generation, which aims to produce high-quality images
closely aligned with user-specified textual prompts. The state-of-the-art
models [7, 11, 37, 1, 3, 30, 31] can generate complex, photorealistic images,
which paves the way to their ubiquitous usage in applications, including web
design, graphics editors, e-commerce and others.

Outstanding image generation ability typically comes at a price of scale,
i.e., state-of-the-art systems require large models, large training datasets,
and vast amounts of GPU resources. At the moment, it is not completely clear
what the trade-offs are between these three aspects of scale and the
generation quality. Moreover, it is unclear if one should prefer large
datasets of random web images over smaller datasets containing only well-
filtered, high-quality ones.

In this paper, we introduce YaART — a cascaded text-to-image diffusion model
enhanced with human preference tuning. YaART is specifically designed to
optimize data and computational resource usage. We thoroughly explore the
impact of model and dataset sizes on generation performance and training
costs. Furthermore, we investigate the quality-vs-quantity trade-off for the
training dataset construction and reveal that training on smaller datasets of
high-fidelity samples provides performance on par with models trained on
larger datasets. Additionally, we show that the model size can be successfully
traded for training time, i.e., smaller models can perform on par with the
larger ones if their training time budget is long enough. Regarding quality,
YaART is consistently preferred in human evaluations over well-known
benchmarks, such as SDXL v1.0, MidJourney v5, Kandinsky v3, and OpenJourney.

To sum up, our main contributions are:

  * •

We introduce YaART, an RLHF-tuned cascaded diffusion model that outperforms
the established models regarding side-by-side human evaluation of image
realism and textual alignment.

  * •

We systematically study the influence of interactions of model and dataset
sizes on the training efficiency and the final model quality. To the best of
our knowledge, this is the first work to extensively investigate the
scalability of diffusion models in convolution-based cascaded schemes using a
reliable human evaluation protocol.

  * •

We investigate the practically important trade-off between data quality and
quantity for pre-training diffusion models. We reveal that competitive results
can be obtained by training on a small amount of high-quality data.

  * •

To the best of our knowledge, we are the first to thoroughly describe RLHF-
tuning of the production-grade diffusion model. We show that RLHF tuning is an
important ingredient of advanced performance; in particular, it substantially
improves image aesthetics and reduces visual defects.

##  2 Related work

Despite the recent progress on large-scale diffusion models, achieving high-
fidelity synthesis of high-resolution images remains a formidable challenge.
Two primary approaches address this issue. The first, employed in models like
DALLE-2 [35] and Imagen [39], termed ‘cascaded’, involves a sequential
pipeline of multiple diffusion models that progressively generate higher-
resolution images. The second approach, known as the Latent Diffusion Model
[37], leverages score matching on the latent space of images generated by a
pre-trained VAE. This method has yielded impressive results and has been
adopted in various subsequent works such as Stable Diffusion XL [31], DALL-E 3
[3], PixArt-$\alpha$ [7], EMU [11].

Traditionally, the denoiser network in diffusion models is implemented as a
variant of the convolutional U-Net model [38] with additional attention
blocks. However, inspired by the success of the transformer architecture [43]
in language modeling, recent research [29, 2, 7, 48] has replaced the
convolutional backbone with a transformer backbone, enabling straightforward
parameter expansion.

In our paper, we adhere primarily to the original approach proposed in DALL-E
2 and Imagen. Specifically, we employ a cascaded scheme of three diffusion
models and a convolutional backbone for denoisers.

####  2.0.1 Alignment of Text-to-Image Models on Human Preferences.

Fine-tuning diffusion models on a clean filtered dataset has been shown to
drastically improve the quality of generated images; however, improving
quality even beyond the level of the available datasets requires other
training methods. Approaches like [26] and [10] suggest generating a dataset
with the current diffusion model, ranking them with scoring models, rejecting
samples with low scores, and fine-tuning diffusion model on this dataset with
MSE loss or reward-weighted MSE loss. In [15], [4], and [27], authors adapt
REINFORCE and PPO for tuning diffusion models. Papers [45] and [44] exploit
DPO for improving quality of generated images, which allows to avoid training
explicit reward functions. Some methods like [10] rely on differentiable
scoring functions while tuning diffusion models.

Diffusion models require multiple steps of inference for generating a single
image; thus, methods like [10] that utilize differentiable rewards and
backpropagate through these rewards and all the iterations of image generation
face a couple of significant problems: vanishing gradients and large GPU
memory requirements for each sample. RL-based approaches like [15], [4], and
[27] require agents to generate trajectories online during training, which
induces additional computational cost. Besides, these methods use on-policy
algorithms, which do not allow them to use images from any sources but the
current model itself. DPO-based approaches like [45] perform on par with D3PO
but have the advantage that they can exploit samples from any other sources.

####  2.0.2 Scaling of Diffusion Models.

The remarkable progress of text-to-image generative models is largely driven
by improvements in scale, where bigger models are trained on larger datasets
for longer schedules. However, the amount of research aimed to understand the
scalability of text-to-image models is still limited. Saharia et al. [39]
investigate the influence of U-Net model size on image quality and conclude
that increasing the model size has little effect, whereas scaling the text
encoder has a greater impact. Caballero et al. [6] build a power scaling law
from experiments on ImageNet 64x64 [12], and LSUN Bedrooms 256x256 [46]
datasets for single-step diffusion models, confirming the potential benefit of
size and compute scaling. Diffusion Transformer (DiT) [29] study the
effectiveness of usage of transformer models in the Latent Diffusion [37]
framework and finds that larger DiT models are more computationally efficient
and that scaling the DiT improves generation quality in all stages of
training.

All mentioned studies mainly rely on FID [18] for image quality estimation,
which has been shown to correlate poorly with human judgments [5, 40].
Moreover, the questions of data quality and quality remain unexplored. In this
work, we study the scalability of the convolutional diffusion model in the
cascade framework based on large-scale study with human annotators. We
investigate the influence of model and dataset sizes on end quality,
specifically considering the data quality-quantity trade-off.

##  3 Approach

Here, we describe our approach to train large-scale diffusion models.

####  3.0.1 High-resolution image synthesis.

We initially opted for the cascaded pixel-based diffusion approach because of
several considerations. Firstly, by leveraging a pre-trained VAE to extract
latent codes for LDM, the quality of the generated images is inherently
constrained by the VAE’s performance, which often has limited capacity.
Although the VAE could be fine-tuned later, such adjustments might not
substantially enhance its overall capability due to inherent capacity
limitations and difficulties with training VAE-based models.

In contrast, cascaded diffusion models allow for important practical
advantages. In cascaded models, the “decoder” function is a stack of two
super-resolution diffusion models that operate independently of the primary
pixel-space generative model. This autonomy allows for the natural
interchangeability of the super-resolution blocks. Furthermore, the diffusion
decoder’s increased capacity and enhanced generative capability significantly
contribute to improving the fidelity of the resultant images.

####  3.0.2 Large high-quality datasets for model pre-training.

One of the most difficult challenges in training a large-scale text-to-image
model is curating a high-quality dataset of text-image pairs. We have
continually refined the process of assembling such a dataset, ultimately
converging on a multi-stage procedure. This approach encompasses separate pre-
filtration stages of image and text filtration and the final stage of
selecting image-text pairs according to their overall fidelity with a focus on
relevance.

####  3.0.3 Boosting relevance and quality via supervised fine-tuning.

Pre-trained text-to-image models often encounter difficulties in image
generation. Exact prompt following, in combination with high visual
attractiveness, is often a challenge. Our strategy involves supervised fine-
tuning using a manually selected subset of data with exceptional aesthetic
quality and detailed descriptions. This process notably enhances prompt
following and relevance, focusing on the latter.

####  3.0.4 Aesthetic refinement through Reinforcement Learning.

In the final stage, we further elevate image aesthetics and quality through
Reinforcement Learning (RL). We directly use the signal provided by human
annotators to improve the aesthetic properties of the produced images.
To sum up, our approach has three main stages. First, we pre-train a diffusion
model on a large-scale dataset, followed by the fine-tuning stage. This model
is further polished via RL-tuning, which incorporates human feedback. Next, we
discuss each component in detail.

###  3.1 Cascaded Diffusion

As discussed above, we base our approach on the framework of cascaded
diffusion models. Specifically, we follow Imagen [39] to design the cascaded
diffusion model. The initial diffusion model in the cascade, which produces
$64\times 64$ images (GEN64), and the first super-resolution model, which
performs up-sampling $64\to 256$ (SR256), follow the architecture of U-Net in
[14, 37, 39, 35]. The GEN64 is conditioned on text via the cross-attention
mechanism as in [37]. The SR256 model is purely convolutional with a single
self-attention in the middle block. Additionally, both models are conditioned
on text via modulation of residual blocks by linear projections of the mean-
pooled text embeddings (as described in [35]). For the final super-resolution
stage $256\to 1024$ (SR1024), we utilize the text-unconditional Efficient
U-Net backbone with the same configuration as the third stage in [39].
Similarly to Imagen, we utilize large pre-trained text encoders for
conditioning. Specifically, we use a CLIP [34]-like proprietary model of size
$\sim 1.3B$ which follows BERT [13]-xlarge architecture. It was pre-trained on
a mixture of public and proprietary datasets of text-image pairs in a setup
similar to OpenAI CLIP models and demonstrated competitive performance on
public benchmarks. Our preliminary smaller-scale experiments verified that
this text encoder provides better quality of generated images according to
human raters than the other choices such as CLIP/OpenCLIP[8]/FlanT5-XXL [9]
encoders. In all our experiments, we use the standard linear noise
scheduler[14] with continuous time; we follow the implementation in [23]. The
timestep $t$ is provided to the model via its respective $LogSNR(t)$, computed
by the scheduler.

###  3.2 Data Selection Strategy

Our proprietary pool of images is comprised of 110B image-text pairs. For each
pair we have a pre-calculated set of scores produced by multiple learned
predictors. These predictors are fully connected classifiers, trained on top
of image and text features, extracted by our proprietary visual and textual
foundation models. For example, we have various classifiers for image
aesthetics and quality, watermark detectors, NSFW detectors, a binary
classifier to check the presence of text within an image, and more.

####  3.2.1 Selecting the best images.

We formulate several stages of filtering for the initial pool of images,
starting from defining the Image Score as a linear combination of the image-
based predictors with weights tuned on the SAC [32] dataset. By visual
inspection in the preliminary experiments, we confirm that the visual appeal
of the images is correlated with the Image Score predictions. We decided to
keep only the top $\nicefrac{{1}}{{3}}$ of the images because most of the
images below this threshold had low visual attractiveness. Additionally, we
remove NSFW images utilizing our proprietary classifiers and take ones with a
size in the range of [512, 10240] pixels and an aspect ratio in the range of
[0.5, 2]. These values were chosen to filter out images with unusual sizes
that may increase computational and engineering burden on our image processing
pipeline or have undesirable content. Then, we filter out images that do not
pass thresholds for the image quality classifiers, specifically those that
contain high levels of noise, blurriness, watermarks, or a checkered
background and were subjected to a high degree of compression. Finally, we
manually select thresholds for aesthetics classifiers, trained on AVA [28] and
TAD66k [16] datasets. We also explicitly control background monotonicity by
splitting the dataset into two parts and allowing only 10% of the samples
during training to have a monotonic background. The motivation is to emphasize
images rich with details during training while retaining the ability to
generate mono backgrounds if required.

Our SR1024 model does not have a condition on the texts. Hence, we obtain the
final version of the 180M Main SR dataset by further tightening the filtering
conditions at this stage, especially those related to image quality.

####  3.2.2 Text filtration.

First, we focus on English texts only by recognizing text language with our
proprietary language detector. Next, we manually annotate a random subset of
4.8K texts so that each raw input text receives its cleaned-up version or an
empty mark if deemed unsuitable for training. Finally, we compile the training
dataset from the refined versions of the texts, filtering out empty text
entries. We then fine-tune our proprietary 180M Language Model on that dataset
and use its prediction as a factor of text quality.

####  3.2.3 Combining All Together.

Previously, we discussed the image filtering process and computation of
various image and text factors. At this point, we have a pool of 2.3B of
relatively poor-quality data, which we aim to filter out further. For that, we
manually label 66K image-text pairs for their visual attractiveness, textual
descriptiveness, and relevance on a scale from 1 to 3. Unlike the more
frequently used Likert scale, the simplified scale balances the
descriptiveness of the score with the noise in the votes of the assessors.
After that, we train a CatBoost [33] model on a set of 56 factors, among which
there are 6 variations of CLIP scores, 38 text-only, and 12 image-only
factors, to obtain Sample Fidelity Classifier (SFC).

All image-text pairs from the previous filtration stage are sorted according
to the SFC model prediction, and the top pairs are selected as the final
training pool. We select a threshold so the final dataset has 300M images with
a non-monotonic background. To obtain the final 330M Main Gen dataset for pre-
training, we also add a random sample of 30M pairs containing images with
monotonic backgrounds.

###  3.3 Model Training

####  3.3.1 Pre-training.

The model is initially pre-trained on a dataset collected as described above.
The GEN64 model with 2.3B parameters is trained on $160$ A100-80GB GPUs with
the total batch size being $4800$ for $1.1\cdot 10^{6}$ training iterations.
The model is trained with automated mixed-precision using ZeRO [36]. We use
Adam optimizer with learning rate being $10^{-4}$ and
$\beta_{1}=0.9,\beta_{2}=0.98$; we observed that the smaller values of
$\beta_{2}$ enhance training stability in half precision. The SR256 model with
700M parameters is trained similarly on $80$ A100-80GB GPUs with total batch
size being $960$ and reduced learning rate $5\times 10^{-5}$; we train it for
$1.5\times 10^{6}$ iterations. The SR1024 EfficientUNet [39], with 700M
parameters, is trained on 256 $\times$ 256 crops under a similar setup, albeit
with a smaller batch size of $512$. We summarize the main information about
training in Table 1. Detailed model configurations are provided in the
supplementary material.

####  3.3.2 Fine-tuning.

To fine-tune the GEN64 and SR256 models, we manually collected a dataset of
50K samples consisting of high-quality images and relevant descriptive
captions. At this stage, we value relevance more than image attractiveness,
which we align later via RL. We also found that fine-tuning the SR1024 model
on full-resolution images from the 180M Main SR dataset leads to a noticeable
improvement in the fine-grained details and sharpness of the images.

Model | Training Stage | Image Resolution | #Samples | Training Steps | Batch Size | Learning Rate
---|---|---|---|---|---|---
GEN64 | Pre-training | 64x64 | 330M Main Gen | 1100K | 4800 | 1 $\times 10^{-4}$
SR256 | Pre-training | 256x256 | 330M Main Gen | 1500K | 960 | $5\times 10^{-5}$
SR1024 | Pre-training | 1024x1024 (crops) | 180M Main SR | 1000K | 512 | $5\times 10^{-5}$
GEN64 | Fine-tuning | 64x64 | 50K | 40K | 240 | $1\times 10^{-5}$
SR256 | Fine-tuning | 256x256 | 50K | 40K | 96 | $1\times 10^{-5}$
SR1024 | Fine-tuning | 1024x1024 | 180M Main SR | 200K | 122 | $2\times 10^{-6}$
GEN64 | RL-alignment | 64x64 | 300K | 50K | 192 | $1\times 10^{-4}$

Table 1:  A summary of training stages of the YaART model cascade. Note that
Main Gen and Main SR are large-scale datasets designed with an intent to cover
a broad range of concepts while the fine-tune datasets consist of carefully
selected image-text pairs, collected with a desire to boost specific quality
aspects such as text-image relevance and image aesthetics.

####  3.3.3 RL Alignment.

For RL alignment, we exploit the PPO [41] approach with $\epsilon=0.5$ for
loss computation and follow DDPO with importance sampling [4]. We also use a
value function that is trained with MSE loss. However, we do not use KL-
divergence and classifier-free guidance because, in our experiments, they
worsen convergence speed, reward growth, and slow down the generation of
trajectories.

For text-to-image diffusion models, the full probability density
$p_{\theta}\left(x_{t-1}|x_{t},c\right)$ is a product of multiple Gaussian
probability density functions for each pixel and channel. It causes rapid
deviation from the value of $p_{\theta_{old}}\left(x_{t-1}|x_{t},c\right)$,
which leads to a significant loss of precision. To cope with this problem, we
split each sample $x_{t}$ into patches and compute the objective for each
patch independently, followed by averaging over all the patches.

We use three reward models: Relevance-focused OpenCLIP ViT-G/14 [22] and
predictors of image Consistency (lack of defects) and Aesthetics, trained on
our human preference data in a Siamese [25] manner with hinge loss as a fully-
connected head over our in-house image foundation model. We independently
compute loss for each reward, followed by averaging them with weights. We aim
to ensure the rise of Aesthetics and Consistency rewards without significantly
declining the Relevance reward. We experimentally found that weighting,
resulting in the same rewards scale, works well in our setup.

To reduce the latency for saving and reading model updates, we train the
condition encoder of our model and LoRA [21] for the remaining layers. We also
tune the frequency of the updates so they occur often enough, as it is crucial
for the on-policy algorithms. We use 100 steps of p-sample [19] so the agent
produces $100\times\text{batch size}$ samples for the loss optimization.

| Quality Aspects
---|---
YaART 2.3B RL | Beauty | Defects | Alignment | Overall
vs MidJourney v5 | $0.58\pm 0.03$ | $0.49\pm 0.01$ |  $0.52\pm 0.01$ | $0.51\pm 0.01$
vs SDXL | $0.75\pm 0.03$ | $0.68\pm 0.02$ |  $0.53\pm 0.02$ | $\textbf{0.77}\pm\textbf{0.01}$
vs Kandinsky v3 | $0.69\pm 0.07$ | $0.69\pm 0.03$ |  $0.46\pm 0.03$ | $\textbf{0.72}\pm\textbf{0.04}$
vs OpenJourney | $0.83\pm 0.09$ | $0.75\pm 0.04$ |  $0.71\pm 0.01$ | $\textbf{0.89}\pm\textbf{0.01}$

Table 2: We compare YaART with state-of-the-art based on three main evaluation
criteria and overall human preferences. Bold denotes statistically significant
difference in overall models’ quality.

###  3.4 YaART

Following the pre-training, fine-tuning, and RL alignment pipeline, we obtain
the final YaART model. While the combination of pre-training and fine-tuning
is relatively common in the modern diffusion training pipelines, we emphasize
the importance of the RL alignment stage, which pushes generation quality even
further. In particular, Figure 2 shows that throughout this stage, the model
learns to generate more aesthetic and consistent (less defective) images while
preserving image-text relevance tuned on the fine-tuning stage.

We compare YaART with MidJourney v5, Stable Diffusion XL with the refiner
module, Kandinsky v3, and OpenJourney. Table 2 demonstrates that YaART
performs on par with MidJourney v5 and is generally preferred over other
counterparts.

Figure 2: An evolution of rewards (left) leads to an increase of human
preference rate (right) throughout the RL alignment stage.

##  4 Experiments

In this section, we first describe the model evaluation approach used in our
experiments. After that, we report the results of a systematic analysis that
provides an in-depth understanding of training dynamics for different choices
of data and model sizes, as well as the trade-off between data quality and
quantity. Finally, we discuss the relationship between the quality of pre-
trained models and their fine-tunes.

###  4.1 Evaluation Setting

####  4.1.1 Prompts.

For evaluation, we use two prompt sets: our practice-driven YaBasket-300
(Figure 3) and DrawBench [39] – a de-facto standard in the field of text-to-
image models. YaBasket complements public benchmarks with several important
practical use cases that have not been previously covered. Specifically, we
enrich small subsets of PartiPrompts [47] and Winoground [42], placed into the
Common Sense category, with popular requests from users of image generation
models and products-driven prompts. For completeness, we perform comparisons
with established reference models in 3.4 using both prompt sets, while all
remaining sections use only the YaBasket set. We publish the prompts from
YaBasket for their subsequent usage in the
community\dagger\dagger$\dagger$Prompts and additional information are
available on the YaART project page..

####  4.1.2 Metric.

Our primary evaluation metric is the Side-by-side (SbS) comparisons between
two images generated from the same prompt. To compute SbS, we use a
proprietary crowdsource platform with non-expert assessors. Before
participating in the labeling process, each candidate assessor must pass an
exam and achieve at least an 80% accuracy rate among a pre-defined set of 20
assignments. After that, we ask assessors to select one of the two images
shown side-by-side based on three evaluation criteria placed in order of their
importance:

  1. 1.

Defectiveness represents the number of visual artifacts and inconsistencies on
an image. Any distortions of objects, limbs, faces, and muzzles belong to this
aspect.

  2. 2.

Image-text Relevance reflects prompt following, including the correctness of
counting, sizes, and positions of objects, their relation, and properties
explicitly stated in the provided textual description.

  3. 3.

Aesthetic Quality is responsible for image style and attractiveness. It
includes color choice and balance, the beauty of the background and
environment, and following the basic rules of photo composition, such as
centering and the rule of thirds.

If the images are indistinguishable based on all three criteria, they are said
to have equal quality, and neither of the images gets a vote. Each pair of
images is labeled by three assessors, after which an image with more votes
gets a point. We conclude that one set of images is better than another by
contesting a null hypothesis of their equality by conducting a two-sided
Binomial test with a p-value of 0.05. We do not report the FID score [18]
because of its poor correlation with human judgments, especially for higher-
quality generations. This phenomenon was already discussed in recent papers
(e.g., [24, 31]), and our analysis in the supplementary material confirms
that.

Figure 3: The content of YaBasket. The three major prompt categories (left)
include Products, almost equally split into eight sub-categories (right).

###  4.2 Scaling Experiments

In this section, we aim to answer the following questions: i) How much compute
is required to train a text-to-image model to reach its quality limits? ii) Is
it possible to obtain a higher-quality model by trading its size (by reducing
the number of parameters) for compute time during pre-training? iii) How do
model and dataset sizes affect the model training dynamic?

First, we define a set of experimental GEN64 models of different sizes: [233M,
523M, 929M, 1.45B], scaled only by width in terms of the number of channels in
the convolutional layers. All other details, such as the number of cross-
attention layers, their position, etc., remain as discussed in section 3.1.

Figure 4: Scaling up the convolutional GEN64 model improves training, leading
to higher quality models. Larger models train faster regarding training steps
and GPU hours, leading to better results across different dataset sizes (top
and middle rows). Dataset size weakly influences the model’s end quality
(bottom row).

####  4.2.1 Learning Rate and Batch Size.

Following the best practices from the area of Large Language Models [20], we
first aim to estimate the optimal hyperparameters (learning rate and batch
size) for training models of various sizes. For that, we perform more than 100
short pre-trainings of the experimental GEN64 models with batch sizes
$\in[48,1152]$ and learning rate $\in[1\times 10^{-8},5\times 10^{-4}]$. All
runs are performed on the full 330M Main Gen dataset discussed in sec. 3.2. By
analyzing training results in terms of FID and CLIP score [17] we conclude
that learning rate of $1\times 10^{-4}$ consistently leads to healthy training
dynamics for larger batch sizes across all model sizes.

####  4.2.2 Model Size, Dataset Size and Compute.

Figure 5: Dynamics of side-by-side comparisons of the half-size model with
Stable Diffusion v1.4 [37] (left) and the fully sized model pre-train (right).
Each point shows the mean and standard deviation between three independent
human evaluation experiments. Note the rapid quality growth through the first
few hundred iterations, after which performance reaches a plateau. Given
enough compute, the test model is capable of surpassing the Stable Diffusion
quality, while the performance of the fully sized YaART model remains
unsurpassed even with more compute.

Previously, DiT [29] showed that scaling transformer denoiser in the LDM
framework [37] improves training in terms of FID. However, the question of the
scalability of convolutional models in the cascade diffusion paradigm remains
open. Moreover, it is unclear how dataset size is related to the model size
and whether there is an optimal combination between the two.

To investigate that, we train four experimental GEN64 models defined above on
five dataset scales, resulting in 20 pre-training runs. The experimental
datasets are obtained by uniform sub-sampling of the 330M Main Gen dataset
with the following sampling fractions:
$[1,\nicefrac{{1}}{{2}},\nicefrac{{1}}{{4}},\nicefrac{{1}}{{8}},\nicefrac{{1}}{{16}}]$.
All models are trained with the same learning rate of $1\times 10^{-4}$ and
batch size of 4800, the same as for the 2.3B model pre-training. We leave the
remaining super-resolution models from the YaART cascade unchanged.

During training, each model is evaluated every 50K training steps with six SbS
comparisons: three against the YaART 2.3B pre-trained model and three against
SD v1.4. We perform all evaluations on YaBasket according to the protocol
described in 4.1. Using different random seeds for image generation and random
pools of human assessors allows us to regularize the evaluation procedure,
avoiding potential bias towards one particular type of images.

Our experiments (Figure 4) show that the quality of convolutional models grows
along with the increased number of parameters. Also, scaling up the models
makes training more efficient in terms of number of training steps and GPU
hours required to achieve quality comparable to our two strong baselines:
Stable Diffusion v1.4 [37] and YaART 2.3B pre-training. Remarkably, this
effect persists when models are trained on datasets of different sizes. The
dataset size alone neither dramatically influences training dynamics nor leads
to substantial changes in resulting model quality. In these experiments, the
models trained on the smallest train set ($\nicefrac{{1}}{{16}}$ of the 330M
Main Gen dataset, around 20M image-text pairs) perform about the same as
models trained on the entire 330M Main Gen dataset. This observation suggests
that datasets of hundreds of millions and, sometimes, billions of image-text
pairs may be excessive, and the dataset size alone is a poor predictor of the
resulting model quality.

####  4.2.3 An Exchange of Model Size for Compute.

The common wisdom in deep Learning is that scaling up a model size and
computational resources typically leads to better models. But can we trade one
for another? To answer this question, we continue pre-training the 1.45B model
from the previous section on the 330M Main Gen dataset. Our goal is to train
it for longer than the 2.3B baseline, both in terms of the number of training
steps and GPU hours, to see whether it allows the test model to surpass the
baseline in terms of generation quality.

In Figure 5, showing the training dynamics, we see that the visual quality
relatively quickly saturates, reliably reaching the performance of SD v1.4 and
YaART 2.3B pre-train around 350K and 500K iterations, respectively. After
that, the quality fluctuates for several hundred thousand iterations and even
stagnates in the later stages of training. We conclude that scaling up the
model size is an unavoidable prerequisite for further growth of generation
quality. We also see that model performance grows non-monotonically,
suggesting that generation quality should be periodically evaluated throughout
training.

###  4.3 Data Quality vs Quantity

Figure 6: Training on a small fraction of high-quality data still leads to
visually appealing generations. Notably, using image aesthetics as a sole
criterion for data selection is constantly inferior compared to the SFC model,
taking into account multiple sample quality aspects.

Most practical settings assume an inherent trade-off between data quality and
data quantity; therefore, the vast majority of application-oriented systems
[7, 11, 31, 1] perform some variation of data selection and filtering during
the dataset collection. However, the influence of the balance between data
quality and quantity on the text-to-image model training dynamics remains
unexplored. In this section, we aim to investigate how a gradual decrease in
dataset size in favor of selecting higher-quality samples affects the dynamics
of model pre-training. We base our analysis on a comparison of the two
approaches for sample ranking used for the initial data filtering: the SFC
model based on 56 image-text factors and the Image Aesthetics model.

Starting from the 330M Main Gen dataset, we progressively drop a fraction of
the worst samples according to one of the aforementioned models. We then train
2.3B models for 300K iterations with default hyperparameters. We use SbS
comparisons to compare the resulting models with the baseline trained on the
full dataset for the same 300K iterations.

Figure 6 shows that the sub-sampling dataset with the Image Aesthetics
classifier only monotonically reduces the quality of the final model, matching
with the intuition that larger datasets might lead to better-performing
models. It also suggests that removing image-text pairs only by means of image
aesthetic attractiveness is suboptimal.

At the same time, sub-sampling data according to the SFC model is more
nuanced. Our evaluations suggest that the best results are obtained at 10%
dataset size, 33M samples. Further reduction of the dataset size employing
more strict filtering leads to a drop in generation quality. There is also a
limit on the minimum size of the dataset, as we show that scaling down to 1%
significantly hurts the performance.

###  4.4 Effect of the model performance on fine-tuning quality.

Figure 7: Pre-train quality strongly correlates with fine-tune quality both
for well-trained models performing similarly with YaART 2.3B pre-train (left)
and smaller models from the Scaling Experiments (right).

In previous experiments, we aimed to understand the influence of various
factors on training dynamics and the quality of the model after the pre-
training stage. However, from the practical perspective, one typically aims to
maximize the final quality, i.e., after fine-tuning the pre-trained model.
However, the relation between quality after these two stages of training is
unclear. To investigate that, we fine-tune various quality models from the
Scaling Experiments 4.2 following the procedure described in Section 3.3.2.

Figure 7 summarizes the results of SbS comparisons between pre-trained models
and their fine-tuned versions. We observe a strong correlation between
performances in these two stages of training despite variations in model size,
dataset size, and dataset quality.

##  5 Conclusion

In this paper, we presented YaART — a production-grade diffusion model for
text-to-image generation. We show that the choices of model and dataset sizes,
as well as the quality of the training images, are important and
interconnected degrees of freedom that should be accurately specified for the
optimal exploitation of the available GPU power. Also, we describe the
procedure of tuning via Reinforcement Learning for YaART, thus confirming that
RL is beneficial for production-level diffusion models.

## References

  * [1] Arkhipkin, V., Filatov, A., Vasilev, V., Maltseva, A., Azizov, S., Pavlov, I., Agafonova, J., Kuznetsov, A., Dimitrov, D.: Kandinsky 3.0 technical report. arXiv preprint arXiv:2312.03511 (2023)
  * [2] Bao, F., Nie, S., Xue, K., Cao, Y., Li, C., Su, H., Zhu, J.: All are worth words: A vit backbone for diffusion models. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 22669–22679 (2023)
  * [3] Betker, J., Goh, G., Jing, L., Brooks, T., Wang, J., Li, L., Ouyang, L., Zhuang, J., Lee, J., Guo, Y., et al.: Improving image generation with better captions. Computer Science. https://cdn. openai. com/papers/dall-e-3. pdf 2(3), 8 (2023)
  * [4] Black, K., Janner, M., Du, Y., Kostrikov, I., Levine, S.: Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.13301 (2023)
  * [5] Borji, A.: Pros and cons of gan evaluation measures. Computer vision and image understanding 179, 41–65 (2019)
  * [6] Caballero, E., Gupta, K., Rish, I., Krueger, D.: Broken neural scaling laws. arXiv preprint arXiv:2210.14891 (2022)
  * [7] Chen, J., Yu, J., Ge, C., Yao, L., Xie, E., Wu, Y., Wang, Z., Kwok, J., Luo, P., Lu, H., et al.: Pixart-alpha: Fast training of diffusion transformer for photorealistic text-to-image synthesis. arXiv preprint arXiv:2310.00426 (2023)
  * [8] Cherti, M., Beaumont, R., Wightman, R., Wortsman, M., Ilharco, G., Gordon, C., Schuhmann, C., Schmidt, L., Jitsev, J.: Reproducible scaling laws for contrastive language-image learning. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 2818–2829 (2023)
  * [9] Chung, H.W., Hou, L., Longpre, S., Zoph, B., Tay, Y., Fedus, W., Li, Y., Wang, X., Dehghani, M., Brahma, S., et al.: Scaling instruction-finetuned language models. arXiv preprint arXiv:2210.11416 (2022)
  * [10] Clark, K., Vicol, P., Swersky, K., Fleet, D.J.: Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.17400 (2023)
  * [11] Dai, X., Hou, J., Ma, C.Y., Tsai, S., Wang, J., Wang, R., Zhang, P., Vandenhende, S., Wang, X., Dubey, A., et al.: Emu: Enhancing image generation models using photogenic needles in a haystack. arXiv preprint arXiv:2309.15807 (2023)
  * [12] Deng, J., Dong, W., Socher, R., Li, L.J., Li, K., Fei-Fei, L.: Imagenet: A large-scale hierarchical image database. In: 2009 IEEE conference on computer vision and pattern recognition. pp. 248–255. Ieee (2009)
  * [13] Devlin, J., Chang, M.W., Lee, K., Toutanova, K.: Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 (2018)
  * [14] Dhariwal, P., Nichol, A.: Diffusion models beat gans on image synthesis. Advances in neural information processing systems 34, 8780–8794 (2021)
  * [15] Fan, Y., Watkins, O., Du, Y., Liu, H., Ryu, M., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Lee, K., Lee, K.: Reinforcement learning for fine-tuning text-to-image diffusion models. Advances in Neural Information Processing Systems 36 (2024)
  * [16] He, S., Zhang, Y., Xie, R., Jiang, D., Ming, A.: Rethinking image aesthetics assessment: Models, datasets and benchmarks. In: Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence, IJCAI-22. pp. 942–948 (2022)
  * [17] Hessel, J., Holtzman, A., Forbes, M., Bras, R.L., Choi, Y.: Clipscore: A reference-free evaluation metric for image captioning. arXiv preprint arXiv:2104.08718 (2021)
  * [18] Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., Hochreiter, S.: Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems 30 (2017)
  * [19] Ho, J., Jain, A., Abbeel, P.: Denoising diffusion probabilistic models. In: Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., Lin, H. (eds.) Advances in Neural Information Processing Systems. vol. 33, pp. 6840–6851. Curran Associates, Inc. (2020), https://proceedings.neurips.cc/paper_files/paper/2020/file/4c5bcfec8584af0d967f1ab10179ca4b-Paper.pdf
  * [20] Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D.d.L., Hendricks, L.A., Welbl, J., Clark, A., et al.: Training compute-optimal large language models. arXiv preprint arXiv:2203.15556 (2022)
  * [21] Hu, E.J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., Chen, W.: Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 (2021)
  * [22] Ilharco, G., Wortsman, M., Wightman, R., Gordon, C., Carlini, N., Taori, R., Dave, A., Shankar, V., Namkoong, H., Miller, J., Hajishirzi, H., Farhadi, A., Schmidt, L.: Openclip (Jul 2021). https://doi.org/10.5281/zenodo.5143773, https://doi.org/10.5281/zenodo.5143773, if you use this software, please cite it as below.
  * [23] Kingma, D., Salimans, T., Poole, B., Ho, J.: Variational diffusion models. Advances in neural information processing systems 34, 21696–21707 (2021)
  * [24] Kirstain, Y., Polyak, A., Singer, U., Matiana, S., Penna, J., Levy, O.: Pick-a-pic: An open dataset of user preferences for text-to-image generation. Advances in Neural Information Processing Systems 36 (2024)
  * [25] Koch, G., Zemel, R., Salakhutdinov, R., et al.: Siamese neural networks for one-shot image recognition. In: ICML deep learning workshop. vol. 2. Lille (2015)
  * [26] Lee, K., Liu, H., Ryu, M., Watkins, O., Du, Y., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Gu, S.S.: Aligning text-to-image models using human feedback. arXiv preprint arXiv:2302.12192 (2023)
  * [27] Lee, S.H., Li, Y., Ke, J., Yoo, I., Zhang, H., Yu, J., Wang, Q., Deng, F., Entis, G., He, J., Li, G., Kim, S., Essa, I., Yang, F.: Parrot: Pareto-optimal multi-reward reinforcement learning framework for text-to-image generation (2024)
  * [28] Murray, N., Marchesotti, L., Perronnin, F.: Ava: A large-scale database for aesthetic visual analysis. In: 2012 IEEE conference on computer vision and pattern recognition. pp. 2408–2415. IEEE (2012)
  * [29] Peebles, W., Xie, S.: Scalable diffusion models with transformers. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 4195–4205 (2023)
  * [30] Pernias, P., Rampas, D., Richter, M.L., Pal, C., Aubreville, M.: Würstchen: An efficient architecture for large-scale text-to-image diffusion models. In: The Twelfth International Conference on Learning Representations (2023)
  * [31] Podell, D., English, Z., Lacey, K., Blattmann, A., Dockhorn, T., Müller, J., Penna, J., Rombach, R.: Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952 (2023)
  * [32] Pressman, J.D., Crowson, K., Contributors, S.C.: Simulacra aesthetic captions. Tech. Rep. Version 1.0, Stability AI (2022), url https://github.com/JD-P/simulacra-aesthetic-captions
  * [33] Prokhorenkova, L., Gusev, G., Vorobev, A., Dorogush, A.V., Gulin, A.: Catboost: unbiased boosting with categorical features. Advances in neural information processing systems 31 (2018)
  * [34] Radford, A., Kim, J.W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al.: Learning transferable visual models from natural language supervision. In: International conference on machine learning. pp. 8748–8763. PMLR (2021)
  * [35] Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., Chen, M.: Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 1(2), 3 (2022)
  * [36] Rasley, J., Rajbhandari, S., Ruwase, O., He, Y.: Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In: Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. pp. 3505–3506 (2020)
  * [37] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., Ommer, B.: High-resolution image synthesis with latent diffusion models. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 10684–10695 (2022)
  * [38] Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18. pp. 234–241. Springer (2015)
  * [39] Saharia, C., Chan, W., Saxena, S., Li, L., Whang, J., Denton, E.L., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Salimans, T., et al.: Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems 35, 36479–36494 (2022)
  * [40] Sajjadi, M.S., Bachem, O., Lucic, M., Bousquet, O., Gelly, S.: Assessing generative models via precision and recall. Advances in neural information processing systems 31 (2018)
  * [41] Schulman, J., Wolski, F., Dhariwal, P., Radford, A., Klimov, O.: Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 (2017)
  * [42] Thrush, T., Jiang, R., Bartolo, M., Singh, A., Williams, A., Kiela, D., Ross, C.: Winoground: Probing vision and language models for visio-linguistic compositionality. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 5238–5248 (2022)
  * [43] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, Ł., Polosukhin, I.: Attention is all you need. Advances in neural information processing systems 30 (2017)
  * [44] Wallace, B., Dang, M., Rafailov, R., Zhou, L., Lou, A., Purushwalkam, S., Ermon, S., Xiong, C., Joty, S., Naik, N.: Diffusion model alignment using direct preference optimization (2023)
  * [45] Yang, K., Tao, J., Lyu, J., Ge, C., Chen, J., Li, Q., Shen, W., Zhu, X., Li, X.: Using human feedback to fine-tune diffusion models without any reward model. arXiv preprint arXiv:2311.13231 (2023)
  * [46] Yu, F., Zhang, Y., Song, S., Seff, A., Xiao, J.: Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv:1506.03365 (2015)
  * [47] Yu, J., Xu, Y., Koh, J.Y., Luong, T., Baid, G., Wang, Z., Vasudevan, V., Ku, A., Yang, Y., Ayan, B.K., et al.: Scaling autoregressive models for content-rich text-to-image generation. arXiv preprint arXiv:2206.10789 2(3), 5 (2022)
  * [48] Zheng, H., Nie, W., Vahdat, A., Anandkumar, A.: Fast training of diffusion models with masked transformers. arXiv preprint arXiv:2306.09305 (2023)

## Supplementary Material

Here we show additional results on the correlation between pre-train and fine-
tune performance for quality vs quantity experiments (Figure 8). We make
additional comparisons with state-of-the-art on DrawBench (Table 3) and
provide a detailed description of hyperparameters (Table 4). Finally, we
discuss YaART’s limitations and provide additional examples of generated
images, showcasing its high visual consistency, aesthetic attractiveness, and
strict prompt-following abilities.

Figure 8: Pre-train quality strongly correlates with fine-tune quality when
trained on a fraction of higher-quality data. This effect is consistent across
two data quality estimation model types: Image Aesthetics (left) and SFT
(right).

### Effect of the model performance on fine-tuning for Quality vs Quantity
trade-off

We previously discussed that pre-train quality strongly correlates with fine-
tune quality in terms of our three-aspect side-by-side (SbS) comparisons
regardless of the model size given a full 330M Main Gen dataset for training.
This experiment aims to understand whether the dataset size and quality
introduce any difference. For that, we fine-tune checkpoints from the Quality
vs Quality experiment, i.e., 2.3B GEN64 models trained on samples of the best
image-text pairs according to Image Aesthetics and Sample Fidelity Classifier
(SFC) models.

Figure 8 shows that even models pre-trained on small sample fractions (1%, 3%)
are easily exposed to fine-tuning. However, even though the best pre-trained
model was obtained by training on the Top 10% of data according to the SFT
model, the model trained on the entire dataset still achieves the best fine-
tuning quality. This suggests that having a large pre-train set of moderate-
quality data is a strong prerequisite for achieving high-end quality.

| Quality Aspects
---|---
YaART 2.3B RL | Beauty | Defects | Alignment | Overall
vs MidJourney v5 | $0.58\pm 0.01$ | $0.54\pm 0.01$ |  $0.52\pm 0.02$ | $\textbf{0.55}\pm\textbf{0.03}$
vs SDXL | $0.78\pm 0.04$ | $0.76\pm 0.03$ |  $0.54\pm 0.02$ | $\textbf{0.82}\pm\textbf{0.02}$
vs Kandinsky v3 | $0.61\pm 0.03$ | $0.70\pm 0.02$ |  $0.48\pm 0.02$ | $\textbf{0.73}\pm\textbf{0.02}$
vs OpenJourney | $0.86\pm 0.06$ | $0.80\pm 0.01$ |  $0.82\pm 0.05$ | $\textbf{0.94}\pm\textbf{0.01}$

Table 3: We compare YaART with state-of-the-art based on three main evaluation
criteria and overall human preferences on DrawBench. Bold denotes a
statistically significant overall model quality difference.

## Limitations

| 64 | 64 $\rightarrow$ 256 | 256 $\rightarrow$ 1024
---|---|---|---
Noise Schedule | _linear_ | _linear_ | _linear_
Sampling Steps | 32 | 32 | 32
Crop Fraction | 1 | 1 | $\nicefrac{{1}}{{4}}$
Model Size | 2.3B | 700M | 700M
Channels | 448 | 320 | 128
Depth | 4 | 4 | 5
Channels Multiplier | $[1,2,3,4]$ | $[1,2,3,4]$ | $[1,2,2,4,8]$
Heads Channels | 64 | 64 | -
Attention Resolution |  $[2,4,8]$ \+ middle | middle | -
Text Encoder Context | 128 | 128 | uncond
Text Encoder Width | 1536 | 1536 | uncond
Dropout | 0 | 0 | 0
Weight Decay | 0 | 0 | 0
Batch Size | 4800 | 960 | 512
Iterations | $1.1\times 10^{6}$ | $1.5\times 10^{6}$ | $1.5\times 10^{6}$
Learning rate | $1\times 10^{-4}$ | $5\times 10^{-5}$ | $5\times 10^{-5}$
Adam $\beta_{1}$ | 0.9 | 0.9 | 0.9
Adam $\beta_{2}$ | 0.99 | 0.99 | 0.99
EMA Decay | 0.9999 | 0.9999 | 0.9999

Table 4:  Hyperparameters for the models in the YaART cascade.

Contrary to the mainstream approach to training text-to-image diffusion models
as LDMs, we have chosen the cascaded diffusion variant. Although this approach
demands more computational resources at runtime, its interactive nature
presents significant benefits: we can provide a 256x256 image that users can
further filter and refine. Moreover, we apply the super-resolution stage only
to a selected fraction of these images. Despite all the recent progress in
this area, modern diffusion models still require substantial human
supervision. This typically involves iterative prompt modification, sampling
parameter tuning (or even adjusting the random seed), and a post-filtering
process for the generated images.

We have also applied automatic filters to eliminate images containing text
from our dataset. This decision stems from our belief that the quality of text
generation currently falls short of practical standards. Interestingly,
despite these limitations, our model has successfully learned to generate
reasonable characters and words from the few text-containing images in our
training dataset.