\section{Mining Docstrings}

\subsection{Mining Argument Schemas}
Sphinx and numpydoc pre-parse the documentation
into a list of argument tuples of the form
\python{(name, short_desc, long_desc)}.
The CNL parser uses the grammar in Figure~\ref{fig:type_grammar_full} to
extract per-argument schemas from the \python{short_desc} fields.

The start symbol of the grammar, \bnf{start}, splits the docstring into
three parts:
the type, or a sequence of possible types (encoded in JSON using the \json{anyOf} keyword); an optional flag; and the default value.

\subsection{Mining Constraints}

For inter-argument constraints, the CNL parser first extracts complete sentences from the long description.
Next, it uses a set of regular expression rules to flag possible candidate constraints, for example, sentences containing the word `only'.
On each candidate it discovers, the CNL parser uses the grammar in \Cref{fig:constraint_grammar_full}.

Unfortunately, there is great variety in how docstrings express conditional hyperparameter constraints.
Since our CNL grammar does not anticipate the syntax of this example, it cannot fully extract this constraint.
However, it does detect the presence of some constraint, and it puts a placeholder into the raw JSON Schema with a \json{TODO} that a human domain expert can then fill in later to further enhance it.

\begin{figure}
\begin{lstlisting}[language=bnf]
start    ::= seq optional default ("." | ",")?
seq      ::= (type ","?)+ ("or " type)?

type     ::= "int" | "integer" | "float" | "double" | "boolean" | "bool" | "string" | "str"
           | "None" | "Ignored"  | "callable" | "dict" | "type"
           | obj | array | enum

optional ::= (", optional")?

default  ::= ","? ("default" ("=" | ":")?  val
           | "(" "default" ("=" | ":")? val ")"
           | val "by " "default"
           | "or " val "(" "default" ")")?

obj      ::= "object" | "RandomState" "instance" | "returns an instance of self"

array    ::= atype (shape)?
atype    ::= "list" | "array" | "tuple" | "array_like" | "array-like"
           | "numpy" "array" | "sparse" "matrix" | "scipy.sparse" |  "scipy" "sparse"
           | "{"? atype ("or " | ",") atype "}"?
shape    ::= ","? "of "? ("shape" | "size")? "="? vtuple ("or " shape)?

enum    ::= "{" val (","? "or "? ("an " | "a ")? val)* "}"
           | ("string" | "str") ","? enum
           | "["? val ("|" val)+ "]"?

vtuple   ::= ("(" | "[") val ("," val)* ","? ("]" | ")") | "None"

val      ::= NAME | NUMBER
\end{lstlisting}
\caption{\label{fig:type_grammar_full}CNL grammar for parsing per-argument schemas.}
\end{figure}

\begin{figure}
\begin{lstlisting}[language=bnf]
start    ::= only when cond

only     ::= "only" ("used" | "effective" | "compatible" | "significant" | "available" | "applies")?
when     ::= "when" | "if" | "with" | "in" | "for"

cond     ::= atom | cond ("and" | "or") cond
atom     ::= NAME compare seq | "the"? seq NAME ("is used")?
seq      ::= val (("," val)* ","? ("and" | "or") val)?

compare  ::= "==" | "=" | ">" | "<" | ">=" | ">=" | "is set to" | "is"
val      ::= NUMBER | NAME
\end{lstlisting}
\caption{\label{fig:constraint_grammar_full}CNL grammar for parsing inter-argument constraints.}
\end{figure}

\section{Results}

This section measures the effectiveness of our extractor tool on two
datasets: complete and curated. The \emph{complete dataset} comprises
119\xspace operators from three different libraries: 115\xspace from
scikit-learn~\cite{buitinck_et_al_2013}, 2 from
XGBoost~\cite{chen_guestrin_2016}, and 2 from
LightGBM~\cite{ke_et_al_2017}.
For XGBoost and LightGBM we considered both the regressor and the classifier.
For scikit-learn we filter the classes revealed by \python{sklearn.utils.testing.all_estimators()} to obtain estimators and transformers.
We exclude classes that are abstract, or that are meta-estimators.
Further we examine their method resolution order, confirm the existence of \python{fit}, and \python{predict} or \python{transform}, and confirm their signatures.
Finally, we exclude some classes known to be deprecated, e.g., \python{Imputer}, and some known to be intended only to be used by other classes, e.g., \python{ExtraTree}.

The \emph{curated dataset} comprises 42\xspace operators with manually curated schemas: 38\xspace
from scikit-learn, 2 from XGBoost, and 2 from LightGBM.  Whereas the
complete dataset allows us to evaluate the robustness and coverage of
the tool overall, the curated dataset allows us to compare the
extracted schemas against a ground truth for reference.

Finally, we used the generated schemas to find three-steps pipelines for 15 OpenML datasets with Lale and compare the results with  with auto-sklearn.

\subsection{Complete Dataset Evaluation}\label{apx:complete}

Table~\ref{table:global_results} presents the results of the extractor
executed on the complete dataset.  For each library, it reports the
number of extracted arguments, types, default values, ranges, and constraints.
In addition, Table~\ref{table:global_results} explicitly shows the contributions of the CNL parser and the schema refiner for each category to illustrate the advantages of combining the two strategies.

Overall, we were able to mine $94\%$ of the~$1,867$ argument types (including the input/output schemas of the \python{fit}, \python{transform}, and \python{predict} methods of all the operators).
Missing argument types mostly come from unsupported or under-specified data structures, e.g., \python{object}, \python{dict}, or \python{callable}.
Even when mined, it is not clear how an AI automation tool would
instantiate such arguments during search.

We extracted a default value for $64\%$ of the arguments but default values are not always relevant, e.g., for the input/output type of the \python{fit} or \python{predict} method.
Table~\ref{table:global_results} shows that, even if default values
are often documented, the schema refiner can extract a lot of
additional information, e.g., via dynamic analysis.

Since ranges are not consistently documented, the range analysis is solely based on the schema refiner.
We found a valid range for~$50\%$ of the~$792$ relevant arguments, i.e., numeric arguments (\json{integer} or \json{number}) or \json{string} arguments used to captures \json{enum} values.
However, ranges are not required for all of these arguments.
In fact we observe that, even in the curated schemas, ranges are only defined for a few arguments to produce valid search spaces for hyperparameter search tools.

Finally, we detected~$120$ constraints but only~$43$ were converted into valid JSON Schema. The remaining~$77$ generate \json{TODO} warnings that can be manually inspected.

\subsection{Curated Dataset Evaluation}\label{apx:curated}

Next, we compare the result of the extractor with a set of manually curated schemas.
Results are presented in Table~\ref{table:curated_results}.
Compared to Table~\ref{table:global_results} we focus on the
arguments to the operator's constructor (\python{__init__}),
leaving aside the arguments and return values of its \python{fit},
\python{predict}, or \python{transform} methods
For each category we report the reference number, the generated
number, and the number of matches between the two sets. We also report
the corresponding precision, recall, and $F_1$ score.

The extractor is able to detect all the~$452$ arguments, which indicates that all the arguments are consistently documented across the three libraries.
In the curated set, all arguments have a correct type and a default value.
The extractor correctly mined the type for~$81\%$ of the arguments and the default value for~$97\%$.
The extractor also found a valid range for~$81\%$ of the~$103$ defined ranges in the curated set.
We categorized a range as valid if the extractor returns an interval
included in the range defined in the curated schema.
Compared to the complete dataset, these counts are relatively low because the curated dataset includes the XGBoost and LightGBM operators that are both more complex and less documented.

Mismatches between extracted and curated default values are due to values that cannot be represented in JSON Schemas: the default value of the \python{missing} argument of the XGBoost operators and the \python{missing_values} of \python{SimpleImputer} is \python{nan}, which is replaced by \python{None} in the curated schemas.
Additionally, our extractor found inconsistencies between the documentation and the code.
For instance, the documentation for the \python{categories} argument of \python{OneHotEncoder} is:
\begin{lstlisting}[language=desc, numbers=none]
categories : 'auto' or a list of lists/arrays of values, \
  default='auto'.
\end{lstlisting}
but the \python{OneHotEncoder.__init__} method gives a default value
\python{categories=None}.

Arguments types are often complex \emph{union} types allowing multiple choices for one argument. For example, \python{n_jobs} can often be either \python{None} or an \json{integer}.
To further investigate the discrepancies between generated and curated types, we analyzed the numbers of \emph{type values} and \emph{enum values}.
The \emph{type values} analysis reports the number of \emph{terminal} types found in the schemas, i.e, \json{boolean}, \json{integer}, \json{number}, \json{string}, or \json{enum}.
The \emph{enum values} analysis reports the number of members found in each \json{enum} list.

\begin{center}
\begin{tabular}{l r r r r r r}
& reference & generated & match &precision & recall & $F_1$\\
\toprule
type values & 631 & 611 & 525 & 0.86 & 0.83 & 0.85\\
enum values & 426 & 269 & 238 & 0.88 & 0.56 & 0.68\\
\bottomrule
\end{tabular}
\end{center}

We observe that for both analyses the precision is relatively high:
$86\%$ for type values and $88\%$ for enum values, which suggests that
extracted data are mostly correct.
However, for enum values the generated number is significantly lower than the curated number: the extractor only found $56\%$ of the enum values in the curated dataset.

This is mostly due to arguments that are under-specified as
\python{string} in the documentation and for which the schema
refiner can not find a suitable enumeration.
For example, the documentation of the \python{criterion} argument
of \python{GradientBoostingClassifier} is:
\begin{lstlisting}[language=desc, numbers=none]
criterion: string, optional (default="friedman_mse")
\end{lstlisting}
but the valid enumeration is
\json{['friedman_mse', 'mse', 'mae']}.
These under-specifications are relatively common when the enum value is only one possible choice in a complex union type.
For example, the documentation of the \python{max_features} argument
of \python{DecisionTreeClassifier} is:
\begin{lstlisting}[language=desc, numbers=none]
max_features: int, float, string or None, optional \
  (default=None)
\end{lstlisting}
but again the valid enumeration is \json{['auto', 'sqrt', 'log2']}.

Finally, the extractor detected~$50$ of the~$65$ constraints of the curated set.
Among the detected constraints, $20$ are converted into valid JSON Schema and~$90\%$ of these match the curated schemas.
The mismatches are due to complex constraints that are merged in the curated schemas.
For instance, the description of the \python{power_t} argument of \python{MLPClassifier} contains:
\begin{lstlisting}[language=desc, numbers=none]
It is used in updating effective learning rate when
the learning_rate is set to 'invscaling'.
Only used when solver='sgd'.
\end{lstlisting}
which is captured in the curated schemas as: "\python{power_t} can only differ from its default value $0.5$ if \python{solver == 'sgd'} and \python{learning_rate == 'invscaling'}", or in JSON Schema:
\begin{lstlisting}[language=json, numbers=none]
'anyOf': [
    {'type': 'object',
     'properties': {'power_t': {'enum': [0.5]}}},
    {'type': 'object',
     'properties': {
        'learning_rate': {'enum': ['invscaling']},
        'solver': {'enum': ['sgd']}}}]}
\end{lstlisting}
but the extractor is only able to extract the second condition on the solver (\python{solver == 'sgd'}).

The results for the constraints show that, even if we are able to flag
most of the constraints, our CNL parser is not the best tool to
extract meaningful information from the constraint candidate.
The language used to express the constraints is far less constrained
than the one used for the type description.  An obvious direction for
future work is thus to try classic natural language understanding
techniques.

\subsection{AutoML Pipelines}\label{apx:autoai}

The selected datasets comprise 5 simple classification tasks (test accuracy~$>90\%$ in all our experiments) and 10 relatively complex tasks (test accuracy~$< 90\%$).
For all the tasks we start from the same three-step pipeline with both generated --- lale-gen in \Cref{tab:accuracy_openml} --- and curated schema --- lale-cur in \Cref{tab:accuracy_openml}.
For comparison, we used auto-sklearn~\cite{feurer_et_al_2015} --- autoskl in \Cref{tab:accuracy_openml} --- as a baseline.
All tasks were configured with the same resource constraints: one hour
of optimization time and a timeout of six minutes per trial.
Lale then uses the search spaces defined in the schemas, the topology of the pipeline, and off-the-self optimizers such as Hyperopt~\cite{komer_bergstra_eliasmith_2014}, to find the best candidate.

\begin{lstlisting}[language=python]
preprocessors       = [ NoOp, MinMaxScaler, StandardScaler, Normalizer, RobustScaler]
features_extractors = [ NoOp, PCA, PolynomialFeatures, Nystroem]
classifiers         = [ GaussianNB, GradientBoostingClassifier, KNeighborsClassifier,
                        RandomForestClassifier, ExtraTreesClassifier,
                        QuadraticDiscriminantAnalysis,  PassiveAggressiveClassifier,
                        DecisionTreeClassifier, LogisticRegression,  XGBClassifier,
                        LGBMClassifier, SVC ]

lale_pipe = make_pipeline( make_choice(*preprocessors),
                           make_choice(*features_extractors),
                           make_choice(*classifiers) )
\end{lstlisting}

For each experiment, we used a $66\%-33\%$ validation-test split, and a 5-fold cross validation on the validation split during optimization.
Experiments were run on a 32 cores (2.0GHz) virtual machine with 128GB memory.
\Cref{tab:accuracy_openml} shows the accuracy results (mean and standard deviation across 5 independent runs) and the number of runs for each experiments (where ``ok'' indicates a successful run and ``ko'' indicates an aborted run).

We observe that Lale with our auto-generated schemas achieves accuracies (88.1\%) that are comparable to auto-sklearn (88.0\% with auto) and Lale with curated schemas (88.5\%).
However, the number of aborted runs show that side-constraints play a key role during the optimization process.
Only 4.0\% of the runs using curated schemas were aborted, compared to 21.7\% with generated schemas (and 6.7\% for auto-sklearn).

\input{table_autoai_results}

\section{Conclusion}\label{sec:conclusion}

This paper presents a tool that mines Python docstrings of ML
libraries to extract hyperparameter schemas for HPO. The extracted
schemas include names, types, defaults, and descriptions, ranges and
distributions for continuous hyperparameters, enumerations of
constants for categorical hyperparameters, and constraints for
conditional hyperparameters.

\section{Mining Docstrings}\label{sec:docstrings}

This section describes the CNL Parser component of the overview
diagram in Figure~\ref{fig:overview}.  A CNL (controlled natural
language) is a natural language (e.g., English) with some
amount of structure (e.g., the input format for Sphinx and numpydoc,
adopted by scikit-learn and other ML libraries).
The CNL parser starts by reading the docstring from the Python file,
such as the one on the left of \Cref{fig:lr_hp}. It uses Sphinx
and numpydoc to extract the docstring of methods \python{__init__}
(class constructor), \python{fit} (for training), and
\python{predict} or \python{transform} (for using an operator
after training) of the operator class.
Sphinx and numpydoc pre-parse this information
into a list of argument tuples of the form
\python{(name, short_desc, long_desc)} as well as
descriptions of the return values of the methods.
Given the list of argument tuples, the
CNL parser uses two hand-crafted CNL grammars to extract per-argument
schemas and inter-argument constraints, respectively.

\paragraph{Mining Argument Schemas.}
The CNL parser uses a grammar (see \Cref{fig:type_grammar_full} in the appendix) to
extract per-argument schemas from the \python{short_desc} fields and inter-argument constraints that appear in the \python{long_desc}.
The example of \Cref{fig:lr_hp} illustrates the main difficulties for mining and
extraction.  The parser needs to ignore noise such as whitespace,
string quotes, or the trailing backslash~(\textbackslash).
In addition, even within the
same operator, the documentation uses different ways to express the
same thing, e.g., enumerating values in curly braces \lstinline!{...}!
vs.\ using \texttt{or}.
Overall, these difficulties arise from the `N' in
CNL: docstrings use natural language.
Noise is easy to handle using filtering during
lexical analysis. For the other difficulties, our grammar takes
advantage of the `C' in CNL: docstrings use controlled language to the
extent that they follow the conventions encouraged by Sphinx and
numpydoc.
The grammar thus specifies multiple syntactic alternatives to capture different ways to express the
same thing (e.g., \bnf{enum} or \bnf{default}).

\paragraph{Mining Constraints.}
For inter-argument constraints, the CNL parser first extracts complete
sentences from the long description. Next, it uses regular
expressions to flag possible candidate constraints, for example,
sentences containing the word `only'. Then it parses each candidate
using a grammar that captures common patterns (see
\Cref{fig:constraint_grammar_full} in the appendix).
Unfortunately, there is great variety in how docstrings express conditional
hyperparameter constraints.
Our grammar is only a first attempt to extract meaningful information.
When the CNL parser fails to parse a sentence flagged as a potential constraint, it puts a
placeholder into the schema with a \json{TODO}
that a human can fill in later.
Having mined both per-argument schemas and inter-argument constraints,
the last step of the CNL parser is to assemble all the pieces into a
single raw schema.
The resulting JSON Schema is machine-readable and captures the
information in a format suitable both for validation and for search.

\section{Introduction}\label{sec:intro}

Machine-learning practitioners use libraries of \emph{operators}:
reusable implementations of estimators (such as logistic regression,
LR) and transformers (such as principal component analysis, PCA).
Training an operator fits its \emph{parameters} (learnable
coefficients such as LR weights or PCA eigenvectors) to a dataset.
Besides parameters, most operators also have \emph{hyperparameters}:
arguments that must be configured before training, such as the choice
of LR solver or the number of PCA components.  Python libraries for
machine learning (ML) such as scikit-learn~\cite{buitinck_et_al_2013} tend
to have good human readable documentation for hyperparameters.
Unfortunately, this documentation is usually not easily machine
readable.
ML practitioners can configure hyperparameters either by hand or by
using an HPO (automated hyperparameter optimization) tool such as
hyperopt-sklearn~\cite{komer_bergstra_eliasmith_2014} or
auto-sklearn~\cite{feurer_et_al_2015},
or the grid search or randomized search from
scikit-learn.  A \emph{hyperparameter schema}
specifies which hyperparameters are categorical and which continuous,
which values or ranges are valid, and conditional hyperparameter
constraints.

\medskip

\begin{wrapfigure}{r}{0.55\textwidth}

  \centerline{\includegraphics[scale=0.45]{figures/overview_2.pdf}}
  \vspace*{-2mm}
  \caption{\label{fig:overview}Overview of our mining approach.}

\end{wrapfigure}

Python has recently emerged as the dominant ML language and many ML
libraries adopt scikit-learn style conventions for interoperability
(including
PyTorch~\cite{skorch},
pandas~\cite{sklearn-pandas},
Spark~\cite{spark-sklearn},
statsmodels~\cite{statsmodels}, and
TensorFlow~\cite{keras-sklearn}).
This
paper proposes and demonstrates an approach for mining hyperparameter schemas from the
Python file implementing an ML operator. The approach, shown in
Figure~\ref{fig:overview}, mines the docstring and refines the
resulting schema via dynamic analysis of the implementation. Using the
Python file as a single source of truth simplifies maintenance when a
new library version adds new features or deprecates old ones.
Furthermore, since the Python file is written by library developers and the
documentation is widely read by library users, it is a reliable source of
truth.

Our approach outputs hyperparameter schemas in \emph{JSON Schema},
which is a type description language for JSON
documents~\cite{pezoa_et_al_2016}.  JSON Schema is widely adopted for
web APIs, cloud management, and document databases, among other
domains, and there are abundant public resources for learning and
using it.  JSON Schema is independent from specific AI automation
tools and recent work has demonstrated that it can be converted to
specifications for popular such tools~\cite{baudart_et_al_2020}.  We found
JSON Schema to have just the right expressiveness for hyperparameters
including categoricals and conditional constraints.  Furthermore, we
found JSON Schema easy to extend with additional meta-data such as
distributions for continuous hyperparameters.

This paper makes the following contributions:
(1)~Mining Python docstrings to extract hyperparameter schemas including constraints.
(2)~Using dynamic analysis to obtain additional information about hyperparameters beyond the docstrings.
(3)~Reconciling hyperparameter metadata into a single machine-readable schema in JSON Schema format.
We evaluate our approach on 119\xspace automatically mined
hyperparameter schemas for  ML operators and 42\xspace hand-curated schemas.
We make both datasets publicly available (\url{https://github.com/IBM/lale/tree/master/lale/lib/}\textsf{\small(autogen$|$sklearn)}).
Overall, we hope this paper contributes towards making HPO tools
easier to use, more reliable, and more effective.

\section{Problem Statement}\label{sec:problem}

This paper is about solving the problem of mining hyperparameter
specifications from a Python docstring and turning them into a JSON
Schema. To make things concrete, \Cref{fig:lr_hp} show an example input and the desired corresponding output of this mining problem.

The left side of \Cref{fig:lr_hp} shows an excerpt of class
\python{sklearn.linear_model.LogisticRegression} with its docstring. A
\emph{docstring} is a string literal that documents a specific class
or function definition. The HTML documentation for scikit-learn and other
popular ML libraries is auto-generated from their docstrings.  For
this to work, the docstrings follow conventions understood by the HTML
generation tool, in this case, the numpydoc
extension~\cite{numpydoc_2008} for Sphinx~\cite{brandl_2008}.  In
other words, docstrings are written in a controlled natural language
(CNL) \cite{kuhn_2014}: controlled, since they follow numpydoc
conventions, and natural, since they are human-readable even before
being converted into HTML. In practice, while docstrings suffice for
HTML generation, they exhibit variability and typos that make schema
extraction non-trivial.

This paper proposes an extractor that converts the docstring not to
HTML but to JSON Schema.  The right side of
\Cref{fig:lr_hp} shows the schema for two categorical arguments
\python{solver} and \python{penalty} and one continuous
arguments \python{C}.
Like most ML libraries, scikit-learn encodes categorical hyperparameters via Python string constants as opposed to Python enums, but only the values
mentioned in the documentation are valid.
The example also contains a conditional hyperparameter constraint:
the value of \python{solver} implies which values are valid for
\python{penalty}. We can express this implication by taking the logically
equivalent form
\mbox{\small$\neg\textit{\textsf{premise}}\vee\textit{\textsf{conclusion}}$}
and using the JSON Schema keywords \json{anyOf} and~\json{not}.

\begin{figure}[t]
\begin{minipage}{0.49\textwidth}
\input{lr_code}
\end{minipage}
\begin{minipage}{0.51\textwidth}
\input{lr_schema}
\end{minipage}
\caption{Simplified excerpt of the scikit-learn code for the LogisticRegression estimator~(left) and the corresponding hyperparameters schema (right).}
\label{fig:lr_hp}
\vspace*{-4mm}
\end{figure}

\section{Refining Mined Meta-Data}\label{sec:refining}

This section describes the Schema Refiner component of the overview
diagram in Figure~\ref{fig:overview}. This component uses dynamic
analysis on the Python code to make additional observations, which
it combines with heuristics and overrides to turn the raw schema
from the CNL Parser into a refined schema for HPO tools.

\paragraph{Dynamic analysis for default values.}
Non-algorithmic defaults complicate the analysis of types
and values.
This occurs, for example, when an argument default is appropriated for
purposes other than parameterization.
To illustrate, it has become
relatively commonplace within scikit-learn to advise users of upcoming
changes in defaults for important arguments by setting the default
value in the constructor's signature to \python{'warn'}
(e.g., Figure~\ref{fig:lr_hp} (left)) to trigger a warning
message.
Our dynamic analysis creates an instance by calling the constructor
\python{__init__()} without passing any explicit arguments, then
calls \python{fit} on the resulting instance --- which might assign the argument to its actual algorithmic default value --- and finally
introspects the instance for these (possibly altered) default values
and their types.

\paragraph{Dynamic analysis via runtime exception testing.}
We use the following techniques to harvest good values and filter bad
values for constructor arguments.

\noindent
\emph{Bad values:}
Defaulting all other arguments, we give a deliberately bad value
to the argument under test and capture the exception.  This
exception text usually reports the bad value which is easily
distinguished in the message.  Frequently, the exception text also
reports valid choices for the argument value, using a range of
syntax that we can parse.

\noindent
\emph{Greedy harvesting:}
We allow argument values that are valid for one operator to be tested
for the same argument name in a different operator.  This occasionally
discovers valid values, particularly for under-documented classes.

\noindent
\emph{Sampling:}
Defaulting all other arguments, an argument's range is sampled for
testing for valid values.  If categorical, all values are tested.
Failed values are filtered out for the class.  The message
received for deliberately false values can help to disambiguate
the complaint in the case where it is not known a priori if
the tested value is good or bad.

\noindent
\emph{Bounds testing:}
With the caveat that some bounds may depend upon
data and the values of other arguments, the bounds of continuous
ranges can be tested individually for validity and
exclusivity, which can be expressed in JSON Schema
via  e.g., \json{'exclusiveMinimum': True}.

\paragraph{Argument Overrides and Relevance to HPO.}
We provide for a dictionary of overrides that allow the automatically
extracted types and ranges to be replaced with user-specified values,
or for the parameter to be excluded from optimization.
E.g., suppressing the \python{'mae'} choice for \python{'criterion'} on tree regressors because of prolonged fit times, or custom bounds for numeric parameters.
The \python{ForOptimizer} suffix indicates that these are not hard constraints.
This step also sets the \python{distribution} field
(e.g.\ \python{loguniform}) and the \python{relevantToOptimizer}
list, omitting irrelevant arguments such as verbosity.

\section{Related Work}\label{sec:related}

The most closely related work is jDoctor, which mines javadoc comments
to extract method pre- and post-conditions~\cite{blasi_et_al_2018}.
The results of mining are similar to schemas in that they can capture
argument ranges and even some constraints. But jDoctor focuses on
Java, whereas we focus on Python code without static type annotations
and with string constants. Furthermore, jDoctor focuses on testing,
whereas we focus on AutoML.

Our schema refiner uses dynamic analysis on Python code to augment the
information extracted from docstrings.  Fuzz
testing, also known as \emph{fuzzing}, is a well-established approach
for finding software defects by generating random
inputs~\cite{miller_fredriksen_so_1990}. While our schema refiner is
inspired by fuzzing, its goal is not to find defects but to
extract schemas.

The primary contribution of this paper is the documentation mining, not
the chosen output format. We could have used different formats to
express hyperparameter schemas.
Python~3 introduces optional type annotations that can be checked
statically~\cite{vanrossum_lehtosalo_langa_2014}. Unfortunately, since
Python~3 types lack intersection types, string constant types, and
conditional constraints, they are less suitable for HPO.
PCS is a file format for specifying hyperparameter schemas for the
SMAC tool~\cite{hutter_ramage_2015}. PCS is well-suited for HPO and
JSON Schema can be converted to PCS~\cite{baudart_et_al_2020}.

\section{Results}\label{sec:results}

This sections measures the effectiveness of our extractor tool with three experiments.
(1)~We mined the schema of 119\xspace operators from three different libraries: 115\xspace from
scikit-learn~\cite{buitinck_et_al_2013}, 2 from
XGBoost~\cite{chen_guestrin_2016}, and 2 from
LightGBM~\cite{ke_et_al_2017}.
(2)~We compared the schemas of 42\xspace operators with manually curated schemas: 38\xspace
from scikit-learn, 2 from XGBoost, and 2 from LightGBM.
(3)~We used the generated schemas to find three-steps pipelines for 15\xspace OpenML datasets with Lale~\cite{baudart_et_al_2020}, an Auto-ML
library that uses hyperparameter search spaces in JSON Schema.

\input{table_global_results}
\paragraph{Complete dataset.}~\\
\Cref{table:global_results} presents the results of the extractor
executed on the complete dataset (see also \Cref{apx:complete}).
For each category, we report the number mined by the CNL parser and the corrections made by the schema refiner ($\textit{parser} + \textit{refiner}$).
For the constraints we report the number of valid constraints and the number of detected constraints ($\textit{valid}/\textit{detected}$).
Overall, we were able to mine $94\%$ of the~$1,758$ argument types (including the input/output schemas of the \python{fit}, \python{transform}, and \python{predict} methods of all the operators).
We extracted a default value for $64\%$ of the arguments but default values are not always relevant, e.g., for the input/output type of the \python{fit} or \python{predict} method.
We found a valid range for~$50\%$ of the~$790$ relevant arguments, i.e., numeric arguments (\json{integer} or \json{number}) or \json{string} arguments used to captures \json{enum} values.
Finally, we detected~$120$ constraints but only~$43$ were converted into valid JSON Schema.

\paragraph{Curated dataset.}
\Cref{table:curated_results} presents the results of the comparison (see also \Cref{apx:curated}).
For this experiment, we focused on the arguments to the operator's constructor.
The extractor correctly mined the type for~$81\%$ of the arguments and the default value for~$97\%$.
\input{table_curated_results}
The extractor also found a valid range for~$81\%$ of the~$103$ defined ranges in the curated set, and~$75\%$ of the distributions.
Finally, the extractor detected~$50$ of the~$65$ constraints of the curated set.
Among the detected cons\-traints, $20$ are converted into valid JSON Schema and~$90\%$ of these match the curated schemas.

\paragraph{Auto-ML pipelines.}
To demonstrate the use of our schemas, we use Lale pipelines of the
form \python{preprocessor >> feature_extractor >> classifier}. Then,
we let AutoML automatically select each step from a predefined set of
operators (see details in \Cref{apx:autoai}) and tune its
hyperparameters based on our extracted schemas.
For comparison, we used auto-sklearn~\cite{feurer_et_al_2015} with the
same resource constraints as a baseline: 1h of optimization time and a
timeout of 6mn per trial.
Note that in this comparison, both the framework and the
hyperparameter schemas differ.
The results show that Lale with our auto-generated schemas achieves
similar accuracies as auto-sklearn, a state-of-the-art tool.


\section{Introduction}
\vspace{-0.3cm}

Generative Adversarial Networks (GANs)~\cite{goodfellow2014generative} have achieved great success in learning high dimensional probability distributions. The capability of GANs advances many important and useful applications, like high quality image generation~\cite{karras2017progressive,brock2018large}, image-to-image translation~\cite{isola2017image,wang2018high}, image editing~\cite{perarnau2016invertible, gu2019mask}, and so on. However, GANs still suffer from the unstable training process.

Many recent works~\cite{radford2015unsupervised,mao2017least,arjovsky2017wasserstein,gulrajani2017improved,miyato2018spectral} focus on stabilizing the training process of GANs. Despite their success, there are still two main issues in current GANs. First, the learned distribution always contains massive low quality samples. We refer to this phenomenon as the \emph{low quality} problem. Second, the learned distribution misses certain areas of real data distribution. We refer to this phenomenon as the \emph{missing modes} (also known as mode dropping) problem. We have carefully studied these two problems and observed that they may be caused by the inaccurate gradient direction from the discriminator.

Ideally, the gradient direction of the discriminator to the generator will make the generated data distribution as close to the real data distribution as possible. However, we found that in practice, the discriminator cannot achieve the theoretical optimal point due to insufficient iteration number or limited fitting ability. In this situation, the gradient direction given by the discriminator is sometimes inaccurate. Previous methods~\cite{gulrajani2017improved,miyato2018spectral} try to solve the inaccurate gradient problem with the Lipschitz constraint for the discriminator. But the Lipschitz constraint will further limit the capability of the discriminator so that it still can not provide accurate gradient for the generated data. More recent work LGAN~\cite{zhou2019lipschitz} proved that penalizing Lipschitz constant
guarantees the gradient for generated data is accurate, but it requires the discriminator need to be optimal before each iteration of generator training, which is impractical for GANs training.

To address these issues, we adopt an auxiliary model to help the generated data get the accurate gradient direction. We propose to build an explicitly prior to the real data distribution for GANs, which we call PriorGANs. Specifically, we choose to adopt the Gaussian Mixture Model (GMM) to establish a probabilistic model for the real data distribution. Equipped with this model, we can address both issues. For the low quality issue, we can apply the model to estimate the probability of each generated image, the probability can also indicate image quality. By encouraging the low quality samples to get a higher probability in the model, low quality samples will get a gradient direction to the real data nearby. Therefore the low quality problem could be solved. We call this the \emph{quality loss}.

For the missing modes issue, we can take full advantage of the GMM model to divide the real data into several groups. Each group may represent a specific mode. Then we can get the densities of the real data in each group. During the training process of GANs, we can estimate the densities of the generated data in each group. If the density of generated data in a group is lower (or higher) than real data, we will increase (or decrease) the sampling frequency of the real data in that group for the next training iteration. We call this the \emph{resampling strategy}. It will help the generative model solve the missing modes problem

The proposed PriorGAN is simple, flexible, and general. Our experiments demonstrate that it is suitable for various types of GANs, and it outperforms state-of-the-art GAN frameworks, including LSGAN, WGAN-GP, SNGAN and StyleGAN, on CIFAR-10, FFHQ, LSUN-cat and LSUN-bird datasets by large margins.

In summary, our key contributions are:

\begin{enumerate}
\vspace{-0.1cm}
    \item We demonstrate the \emph{low quality} and \emph{missing modes} problems in current GANs and demonstrate the cause of these two problems on toy examples.
    \vspace{-0.1cm}
    \item We propose \emph{real data prior} for GANs(PriorGANs), equipped with the prior, we propose the \emph{quality loss} and \emph{resampling strategy} help current GANs solve these problems.
    \vspace{-0.1cm}
    \item Our experiments demonstrate that PriorGANs outperform state-of-the-art GAN frameworks on various datasets by large margins.
    \vspace{-0.1cm}
\end{enumerate}

\iffalse

\fi

\begin{comment}
\begin{figure*}[t]
\centering
 \includegraphics[width=1.0\columnwidth]{problem.pdf}
 \vspace{-0.7cm}
\caption{Two problems in GAN training, the imprecise boundary problem and mode missing problem.}
\label{fig:gan_problem}
\vspace{-0.2em}
\end{figure*}
\end{comment}
\vspace{-0.3cm}
\section{Related Work}
\vspace{-0.3cm}
Over the past years, GANs have achieved significant attention in the deep learning community. Massive works are proposed to improve the performance of GANs. In this section, we will briefly review them. Also, we will review some classical algorithms about modeling data distribution.

\noindent \textbf{Generative Adversarial Nets}.
Early works about GANs try to solve the unstable training process of GANs, many works focus on proposing new loss functions to solve this problem. For example, LSGAN~\cite{mao2017least} replace the original BCE loss with the proposed least square loss. WGAN~\cite{arjovsky2017wasserstein} proposes the Wasserstein distance and weight clipping to achieve the Lipschitz condition. Later work began to notice the Lipschitz condition. Some typical works are as follows:
WGAN-GP~\cite{arjovsky2017wasserstein} further introduce the gradient penalty to achieve the Lipschitz condition. SNGAN~\cite{miyato2018spectral} proposes Spectral Normalization to achieve Lipschitz condition for training GANs. More recent works~\cite{dong2019margingan,NIPS2019_8560,cao2019multi} design task-specific loss functions for GANs. Meanwhile, some works aim to improve the diversity of the GANs, for instance, VEEGAN~\cite{srivastava2017veegan} and CVAE-GAN~\cite{bao2017cvae} proposes to reconstruct the images in the pixel level to avoid the missing modes problems. MS-GAN~\cite{mao2019mode} proposes a loss term that maximizes the ratio of the distance between generated images with respect to the corresponding latent codes. Besides new loss functions, there are also some works~\cite{radford2015unsupervised, karras2017progressive, zhang2018self, karras2019style} focus on using different network structures and training strategies for training. In contrast to these methods, our method proposes to build explicit real data prior to solving these problems.

\noindent \textbf{Build Prior for Data Distribution}.
There are various methods that aim to build the prior for underlying data distribution. For example, Principle Component Analysis (PCA) ~\cite{turk1991face}, Independent Component Analysis (ICA)~\cite{hyvarinen2000independent}, and the Gaussian Mixture Model (GMM)~\cite{xu1996convergence, richardson2018gans}, all assume the data distribution follow a simple assumption. But they have difficulty modeling complex patterns of irregular distributions. Later works, such as the Hidden Markov Model (HMM)~\cite{starner1997real}, Markov Random Field (MRF)~\cite{ranzato2010generating}, and restricted Boltzmann machines (RBMs)~\cite{salakhutdinov2009deep}. limiting their results on texture patches, digital numbers, due to a lack of effective feature representations. Instead of building the model directly on the pixel space, we leverage the successful feature presentation of the CNNs~\cite{He_2016_CVPR,szegedy2016rethinking} and build the model on these feature representations.

\vspace{-0.3cm}
\section{Methods}

Before we introduce our methods, we first give a brief introduction and our findings about current GANs. In the vanilla GAN, the discriminator D and the generator G play the following two-player minimax game with value function $V(D, G)$:

\begin{equation}
\label{eq:minimaxgame-definition}
\min_{G} \max_D V(D, G) = \mathbb{E}_{x \sim p_{\text{r}}(x)}[\log D(x)] + \mathbb{E}_{x \sim G(z)}[\log (1 - D(x))].
\end{equation}

suppose $p_g(x) = G(z)$, for G fixed, the optimal discriminator $D$ for $V(D, G)$ is:
\begin{equation}
\label{eq:optimal-D}
   D^*_G(\bm{x}) = \frac{p_\text{r}(x)}{p_\text{r}(x) + p_g(x)}
\end{equation}

According to the analysis in GAN~\cite{goodfellow2014generative}, the global minimum of $V(D^*_G, G)$ is achieved if and only if $p_g = p_r$. But in the practical training process of GAN, the optimal discriminator $D_G^*(x)$ is usually hard to achieve due to insufficient training iterations or limited network capability. In this situation, the gradient is probably incorrect for reaching the optimal $p_g$.

To demonstrate this, we train a basic GAN~\cite{goodfellow2014generative} on two typical toy distributions. As illustrated in Figure~\ref{fig:toy_vis}, The green points are the real data, the blue points are the generated data, the discriminator is trained based on current real and generated data distribution. The red arrows denote the gradient of the discriminator for each generated data. In Figure~\ref{fig:toy_vis}(a), some generated data are out of the real data distribution. But the gradients of these low quality generated points are not correct and they can't move to the real data distribution. So these samples stay in the situation of \emph{low quality} problem. In Figure~\ref{fig:toy_vis}(c), we notice that the real data distribution has two main regions, but the gradient of the generated points have the same direction to one region of the real data. This seems to be an intrinsic cause of the \emph{missing modes} problem.

\begin{figure*}[t]
\centering
 \includegraphics[width=1.0\columnwidth]{vis.pdf}
 \vspace{-0.3cm}
\caption{An example of \emph{low quality} problem(a,b) and \emph{missing modes} problem(c,d) in GANs training on toy examples, (a) and (c) show the gradient of vanilla GAN, (b) and (d) show the gradient of our proposed PriorGAN.}
\label{fig:toy_vis}
\vspace{-0.3cm}
\end{figure*}

To solve these problems, we propose to establish a prior model on real data to help the discriminator solve the inaccurate gradient problem. Equipped with the prior, we can solve the problem of \emph{low quality} and \emph{missing modes}.
In the next section, we will first introduce how to build real data prior. Then we present how to apply the prior as a loss term in GAN training to improve the quality of generated images. After that, we show how to use this prior to estimate the diversity of the real and generated data distribution. Finally, we introduce the resampling strategy to improve the diversity of the generated samples.

\subsection{Build Prior for Real Data}
Building an explicit prior to the high dimensional data is very difficult and challenging, especially for high-resolution images. This is because the real data distribution is extremely complicated to be captured by a simple prior. Previous method~\cite{richardson2018gans} try to leverage the Gaussian Mixture Model(GMM) model to build a model on the image patch level. However, the image resolution is restricted to be very small and the generated results are really blurry. Instead of directly building the GMM model on the pixel level, we take full advantage of the recent success of deep convolutional networks and build the GMM model on the features extracted by the networks. The deep feature is suitable for building the GMM model from two aspects. First, the deep feature is relatively low dimensional, which helps avoid the overfitting problem. Second, the deep feature is very representative and explicitly represents the content of the image. Denote the input image as $x$, the deep feature extractor function as $F(\cdot)$. We use a pre-trained network as feature extractor and the extracted feature dimension is $d$. Suppose the GMM model has $M$ Gaussian components, the mean vector and covariance matrix for $i$-th Gaussian component are $\mu_i$ and $\Sigma_i$, respectively. Given these parameters, the density function of a Gaussian Mixture Model takes the form
\begin{equation}
    p(x|\lambda) = \sum_{i=1}^{M} w_i \mathcal{N}(F(x)|\mu_i, \Sigma_i),   x \in P_r
    \label{eqn:GMM_probability}
\end{equation}

$P_r$ denotes real image space, $w_i$ is the mixture weights, which satisfy the constraint that $\sum_{i=1}^{M} w_i = 1$. And $\mathcal{N}(F(x)|\mu_i, \Sigma_i)$ is a $d$-variate Gaussian function of the form. The complete Gaussian mixture model is parameterized by the mean vectors, covariance matrices and mixture weights from all component densities. These parameters are collectively represented by the notation, $\lambda = \{w_i, \mu_i, \Sigma_i\}$. To estimate these parameters, we adopt the expectation-maximization (EM) algorithm~\cite{dempster1977maximum} to iteratively update them with the loss function
\begin{equation}
L= \sum_{x \in P_r} -\log p(x|\lambda).
\vspace{-0.2cm}
\end{equation}

\subsection{Improving Quality with Prior}
Given the estimated GMM model, for a generated image $x \in P_g$, we can adopt $p(\mathbf{x}|\lambda)$ in Equation.~\ref{eqn:GMM_probability} to estimate its probability. If its probability is high, it indicates that it is close to some of the real samples and the quality is high, otherwise, it is far from real samples and the quality is poor. So the generator should pay more attention to these poor quality samples. However, previous analysis shows that these poor quality generated samples may get inaccurate gradient from the discriminator. To solve this problem, we propose to utilize the prior to restrict these poor quality samples to get a higher probability in the built GMM model. This will help the poor quality samples to have an accurate gradient direction towards the real data distribution. Formally, given an generated image $x = G(z)$, $G$ is the generator, the loss for the generator is:

\begin{equation}
\label{eqn:quality_score}
\mathcal{L}_{Q} =
\begin{cases}
-\log p(F(G(z))| \lambda) & \textit{if } p(F(G(z))| \lambda) < \theta\\
0 & \textit{otherwise}
\end{cases},
\end{equation}
where $\theta$ is the quality threshold. The quality loss $\mathcal{L}_{Q}$ mainly focuses on the low quality generated results and will give a right gradient direction for these low quality generated samples, as illustrated in Figure~\ref{fig:toy_vis}(b). During the training process of GANs, we give a loss weight $\delta$ to this quality loss and apply it as an auxiliary term to update the generator besides the GAN loss.

\subsection{Improving Diversity with Prior}
Current GANs often suffer from the \emph{missing modes} problem, the learned distribution from GANs may miss some certain regions of the real data distribution. Therefore the key is how we can estimate the diversity of real data distribution and find these missing regions in the learned distribution. One possible solution is to apply the built GMM model. Specifically, if we establish a GMM model to represent a real data distribution, then we can apply the $M$ Gaussian components in the GMM model to divide the real data distribution into $M$ groups and estimate the frequency of the real samples in each group. A sample belonging to which Gaussian component is decided by the minimum distance to the center of each Gaussian component, as we shown in Figure~\ref{fig:prior}. Suppose the number of the real samples on the $i$-th Gaussian component is $n_i^r$, then the frequency of the real samples on the $i$-th Gaussian component is $f_i^r = {n_i^r}/{\sum_i^M{n_i^r}}$.  And we have the vector $[f_1^r, f_2^r,\cdots, f_M^r]$ to represent the diversity of the real samples. Similarly, we apply the same GMM model to get the frequency of generated samples, which is denoted as $[f_1^g, f_2^g,\cdots, f_M^g]$.

\begin{figure*}[t]
\centering
 \includegraphics[width=1.0\columnwidth]{split.pdf}
 \vspace{-0.6cm}
\caption{Visualization of learned GMM-based prior on FFHQ and LSUN-cat dataset, each column are images randomly selected from the same Gaussian component. We can observe that images within the same Gaussian component share the same attributes, like hat, hair, sunglasses for the FFHQ or eyes, color, posture for the LSUN-cat.}
\label{fig:prior}
\vspace{-0.2cm}
\end{figure*}

\begin{figure*}[t]
\centering
 \includegraphics[width=1.0\columnwidth]{resample.pdf}
 \vspace{-0.7cm}
 \caption{illustration of the pipeline of our resampling strategy.}
\label{fig:resample}
\vspace{-0.4cm}
\end{figure*}

The frequency difference between the real and generated samples can indicate whether the missing modes problem occurs. If the learned distribution misses a certain region $S$ of the real data distribution, and $S$ belongs to the $i$-th Gaussian component of the built GMM model, then $i$-th learned data frequency $f_i^g$ will be smaller than the real data frequency $f_i^r$. So the diversity difference between real data distribution and learned distribution can be represented by the frequency distances. Suppose the diversity distance between real data distribution $p_r$ and learned distribution $p_g$ is denoted as
\begin{equation}
    d(p_r,p_g) = [f_1^r, f_2^r,\cdots, f_M^r] - [f_1^g, f_2^g,\cdots, f_M^g].
\end{equation}
So $d(p_r, p_g)$ is a vector with $M$ elements, the $i$-th element is $d_i(p_r,p_g)$. So the overall diversity differences score (DDS) is
\begin{equation}
DDS(p_r, p_g) = \sum_i^M |d_i(p_r,p_g)|.
\end{equation}

An ideal learned distribution should satisfy $DDS(p_r, p_g)=0$. However, the $DDS$ is non-differentiable and cannot be directly used as a loss function. So we optimize it in an indirect way and proposed a resampling strategy. Considering a typical missing modes phenomenon where the real data distribution has two regions $S_1$ and $S_2$, the generated samples only lies around the $S_1$, current discriminator can't force the generated samples around $S_1$ to move to $S_2$ since the discriminator consider the generated samples around $S_1$ are "real enough". To overcome this problem, we propose to change the decision boundary of the discriminator. If we sample fewer samples in $S_1$ and more samples in $S_2$ for training the discriminator, the decision boundary of the discriminator will be changed. The discriminator will encourage the generated samples to move to $S_2$ and solve the \emph{missing modes} problem, as illustrated in Figure~\ref{fig:toy_vis}(d). We call this process the resampling strategy. This strategy will change the real data distribution in the training procedure to help the generator overcome the missing modes problem. Formally, we use a new sampling frequency $f^{r'}_i$ to sample the real samples in the $i$-th Gaussian component for training. $f^{r'}_i$ is decided by the original frequency and the diversity distances:
\begin{equation}
    f^{r'}_i = max(f_{i}^{r} + \alpha d_i(p_r,p_g), 0),
\end{equation}
Then we normalized the sampling frequency, $f^{r'}_i = {f^{r'}_i}/ {\sum_i^M f^{r'}_i}$. Such that the new sampling frequency is between $0$ and $1$, $\alpha$ is a hyper parameter which decides the resample weight. For $f_g^i < f_r^i$, we increase the sampling frequency of region $i$ and decrease it if otherwise. The resampling strategy is also illustrated in Figure~\ref{fig:resample}.

\iffalse
In high dimensions, it's not easy to get $p_r(x)$, $p_g(x)$ for a given sample $x$. so we use our estimated GMM model to get the approximate number. In details, if we split the real data distribution to $M$ subspace, calculate the frequency of real data according to the identity of nearest Gaussian Kernel, we denote the frequency as $\mathbf{F_{r}_i}, i={1,2,..M}$, and the frequency for generated data as $\mathbf{F_{g}_i}, i={1,2,..M}$, these can be seen as
the summation of $p_r$ and $p_g$ in each subspace, and it satisfied:
\begin{equation}
    \sum_{\mathbf{i}=1}^\mathbf{M} \mathbf{F_{r}_i} = \sum_{\mathbf{i}=1}^\mathbf{M} \mathbf{F_{g}_i} =  1
\end{equation}
we denote the difference between these two frequency distribution as $\mathbf{D(r,g)}$:
\begin{equation}
    \mathbf{D(r,g)} = \sum_{\mathbf{i}=1}^\mathbf{M} |\mathbf{F_{r}_i}-\mathbf{F_{g}_i}|
\end{equation}
We named it DIV-score. An ideal generated distribution should satisfy $\mathbf{D(r,g)}=0$. And a low DIV-score is a necessary condition which indicates that generated distribution fits the real data distribution.

But in practical application, the DIV-score might not be very small to some GAN models, this usually indicates modes missing. To narrow the gap between $p_r$ and $p_g$, instead of sampling from real data distribution directly, we aim to balance the sampling procedure to remedy the missing mode.
Specifically, we assume the default sampling frequency for $\mathbf{i}$-th Gaussian component is $\mathbf{F_r_i}$, we update it to $\mathbf{F_n_i}$:
\begin{equation}
    \mathbf{F_{n}^{i}} = \mathbf{F_{r}^{i}} + (\mathbf{F_{r}^{i}} - \mathbf{F_{g}^{i}}) * K,   K>0
\end{equation}

K is a hyper parameter which decides the resample weight. For $F_g^i < F_r^i$, we increase the sampling frequency of region $i$ and decrease it if $F_g^i > F_r^i$.
\fi

Above we introduce our methods for improving quality and diversity with our proposed prior. This prior can be integrated into various frameworks of GANs, such as WGAN-GP~\cite{gulrajani2017improved}, LSGAN~\cite{mao2017least}, SNGAN~\cite{miyato2018spectral}, and even the StyleGAN~\cite{karras2019style}. This flexible combination with existing methods can help these methods get further improvement.

\section{Experiments}

In this part, we conduct experiments on a large variety of generative models trained on different datasets including CIFAR-10~\cite{krizhevsky2009learning}, LSUN~\cite{yu2015lsun}, and FFHQ~\cite{karras2019style} datasets to evaluate the performance of our proposed methods.

\textbf{Datasets.}
CIFAR-10 is a widely used dataset for image generation. We directly apply the original image size $32 \times 32$ for training. LSUN~\cite{yu2015lsun} is widely used for high-resolution image generation, we choose two categories in LSUN for training and evaluation: cat and bird, which we call LSUN-cat and LSUN-bird, respectively. LSUN-cat contains $1.6$ million images and LSUN-bird contains $2.2$ million images, we follow the settings in StyleGAN~\cite{karras2019style} to resize all the images to $256 \times 256$ resolution. FFHQ~\cite{karras2019style} is a recently high-resolution face dataset contains $70000$ images, we resize them to $256 \times 256$ resolution for training, which we call FFHQ-256.

\textbf{Evaluations.}
To evaluate the quality of an image set, as we mentioned earlier, the probability of each image in our GMM model can indicate its quality. So we use the mean of normalized image density in Equation~\ref{eqn:GMM_probability} to evaluate the quality of an image set, calculated as:  $ \frac{1}{K} \sum_{\mathbf{i}=1}^K norm( log ( p(x_i|\lambda)))$, $K$ is the number of images in the image set\footnote[1]{More implementation details please refer to supplementary material.}. We named it Quality Score(QS), higher QS indicates better image quality. And we use our proposed DDS to evaluate the diversity of the generative model, FID~\cite{heusel2017gans} to measure both quality and diversity.
For a fair comparison, we would use $5k$ or $50k$ genearted images to calculate the FID, denoted as FID(5K) and FID(50K) respectively. $10k$ images are used to calculate QS and DDS for the CIFAR-10 dataset. For the FFHQ, LSUN-cat, and LSUN-bird datasets, we use $50k$ images to calculate the result of the FID, QS and DDS.

\textbf{Implementation Details.}
For each baseline, we follow training settings in their experiments. We only add quality loss and the resampling strategy during the training of each baseline. In order to build the prior for each dataset, we apply ResNet101~\cite{He_2016_CVPR} or InceptionV3~\cite{szegedy2016rethinking} to extract the features for all the real images. The features we use are from the last average pooling layer. Without special notes, we use ResNet101 to extract the features.
Then we build the prior for each dataset. Specifically, we set the number of Gaussian component to $35$,$70$, $7$, and $7$, the resample weight $\alpha$ to $3$,$7$,$7$,$7$ for the CIFAR-10, FFHQ, LSUN-cat, and LSUN-bird, respectively, and the quality loss weight $\delta$ set to $0.1$. More details could be found in the supplementary material.

\begin{table}[t]
\centering
  \small{
  \begin{tabular}[t]{c|ccc}

    \hline
    Method & Inception score $\uparrow$ & FID(5k)$\downarrow$ &  FID(50k)$\downarrow$ \\

    \shline
    Real data & 11.24$\pm$.12  & 8.63$\pm$.0 & - \\

    \hline
    DCGAN~\cite{radford2015unsupervised}& 6.64$\pm$.14 & - & - \\
    LR-GANs~\cite{yang2017lr} &  7.17$\pm$.07 & - & - \\
    Warde-Farley et al.~\cite{warde2016improving} & 7.72$\pm$.13 & - & - \\
    WGAN-GP~\cite{gulrajani2017improved} & 7.86$\pm$.08 & 40.2$\pm$.0 & - \\

    LGAN~\cite{zhou2019lipschitz} & 8.03$\pm$.03 & - & 15.64$\pm$.07 \\
    SNGAN~\cite{miyato2018spectral} & 8.22$\pm$.05 & 21.7$\pm$.21 & 13.78$\pm$.11  \\
    $L^{10}$ BWGAN~\cite{adler2018banach} & 8.31$\pm$.07 & - & -  \\
    QAGAN~\cite{NIPS2019_8560} & 8.37$\pm$.04 & - & 13.91$\pm$.11  \\
    WGAN-ALP~\cite{terjek2019virtual} & 8.34$\pm$.06 & - & 12.96$\pm$.35 \\

    PriorSNGAN(\textbf{ours}) & \textbf{8.60$\pm$.19} & \textbf{18.19$\pm$.06} & \textbf{ 11.94$\pm$.18}   \\

    \shline
  \end{tabular}
  \vspace{0.1cm}
    \caption{\label{tab:incepscores}Inception scores and FIDs with unconditional image generation on CIFAR-10. FID$(50k)$ means calculating the score with $50k$ generated images, otherwise with $5k$ generated images.
  }
  \vspace{-0.5cm}
  }
\end{table}

\begin{comment}
\begin{table}[t]
\centering
  \small{
  \begin{tabular}[t]{lcc}
    \toprule
    Method & Inception score $\uparrow$  &  FID(50k)$\downarrow$ \\

    \midrule
    Real data & 11.24$\pm$.12  &  \\
    \midrule
    DCGAN~\cite{radford2015unsupervised}& 6.64$\pm$.14 & \\
    LR-GANs~\cite{yang2017lr} &  7.17$\pm$.07\\
    Warde-Farley et al.~\cite{warde2016improving} & 7.72$\pm$.13 &  \\
    WGAN-GP~\cite{gulrajani2017improved} & 7.86$\pm$.08 & 19.65 \\

    LGAN~\cite{zhou2019lipschitz} & 8.03$\pm$.03 & 15.64$\pm$.07 \\
    SNGAN~\cite{miyato2018spectral} & 8.22$\pm$.05 & 13.78$\pm$.11  \\
    $L^{10}$ BWGAN~\cite{adler2018banach} & 8.31$\pm$.07 & -  \\
    WGAN-ALP~\cite{terjek2019virtual} & 8.34$\pm$.06 & 12.96$\pm$.35 \\

    PriorSNGAN(\textbf{ours}) & \textbf{8.60$\pm$.19} & \textbf{ 11.94$\pm$.18}   \\
    \bottomrule
  \end{tabular}
  \vspace{0.1cm}
    \caption{\label{tab:incepscores}Inception scores and FIDs with unconditional image generation on CIFAR-10. FID$(50k)$ means calculating the score with $50k$ images otherwise with $5k$ generated images.
  }
  }
\end{table}
\end{comment}

\subsection{Comparison with state-of-the-art Methods}
In this section, we compare our results with other state-of-the-art methods. We conduct the experiments on the CIFAR-10 dataset. The basic framework is an unconditional GAN, the generator and discriminator network structure are the same as the ResNet structure adopted in the baseline SNGAN~\cite{miyato2018spectral}. We use $5k$ randomly generated images to calculate the Inception score ten times and report the mean and variance. For the FID, there are two different settings to calculate the result, some methods~\cite{miyato2018spectral, heusel2017gans} use $5k$ generated images and $10k$ real images to calculate, while the others~\cite{parimala2019quality,terjek2019virtual} use $50k$ generated and $50k$ real images to calculate. Using $50k$ generated images can get a better result. For a fair comparison, we report both results and denote them as the FID($5k$) and the FID($50k$). As shown in Table~\ref{tab:incepscores}, our method (PriorSNGAN) outperforms the baseline SNGAN and other state-of-the-art GANs. These results confirm the effectiveness of our proposed real data prior to GAN training.

\begin{table}[t]
\centering
  \begin{tabular}{c|ccc|ccc|ccc}
    \hline
    \multirow{2}{*}{ } &
      \multicolumn{3}{c|}{SNGAN } &
      \multicolumn{3}{c|}{WGAN-GP } &
      \multicolumn{3}{c}{LSGAN } \\
    \cline{2-10}
    &  FID$\downarrow$ &  QS$\uparrow$  & DDS$\downarrow$ & FID$\downarrow$& QS$\uparrow$ & DDS$\downarrow$ &  FID$\downarrow$ &  QS$\uparrow$  & DDS$\downarrow$   \\
    \shline
    Baseline & 13.78 & 0.429  & 0.488 & 19.65 & 0.410  & 0.532 & 34.90 & 0.367  & 0.545 \\

    +Prior & 11.94 & 0.512  & 0.380 & 18.16 & 0.439  & 0.475 & 30.24 & 0.406 & 0.487 \\
    \shline
  \end{tabular}
  \vspace{0.14cm}
\caption{Comparison of FID, QS, and DDS metrics on three different GANs: SNGAN, WGAN-GP and LSGAN.}
\label{table:model_score}
\vspace{-0.5cm}
\end{table}

\subsection{Prior for Various GANs}
Our approach enjoys a high degree of flexibility and can be integrated into various kinds of GAN frameworks. We choose some variants of GAN as baseline frameworks, such as LSGAN~\cite{mao2017least}, WGAN-GP~\cite{gulrajani2017improved}, SNGAN~\cite{miyato2018spectral}, and then integrate our method into these frameworks for comparison. For a fair comparison, we follow the same settings and architecture of GANs and only apply our quality loss and the resampling strategy, which are denoted as +Prior.

We conducted experiments on the CIFAR-10 dataset, Table~\ref{table:model_score} reports a comparison of our method with baseline models. We can observe that our methods outperform the corresponding baseline GAN frameworks in terms of FID, QS, and DDS metrics, which demonstrate the superiority of our proposed framework.

\begin{figure*}[t]
\centering
 \includegraphics[width=1.0\columnwidth]{visualization_results.pdf}
 \vspace{-0.4cm}
\caption{Visualization results of our generated images on FFHQ-256, LSUN-cat and LSUN-bird datasets, our method can generate not only high quality but also diverse images.}
\label{fig:vis_results}
\vspace{-0.1cm}
\end{figure*}

\begin{table}[t]
\centering
  \begin{tabular}{c|ccc|ccc|ccc}
    \hline
    \multirow{2}{*}{ } &
      \multicolumn{3}{c|}{FFHQ-256 } &
      \multicolumn{3}{c|}{LSUN-cat } &
      \multicolumn{3}{c}{LSUN-bird } \\%&

    \cline{2-10}
    &  FID$\downarrow$ &  QS$\uparrow$  & DDS$\downarrow$ & FID$\downarrow$& QS$\uparrow$ & DDS$\downarrow$ &  FID$\downarrow$ &  QS$\uparrow$  & DDS$\downarrow$   \\

    \shline
    StyleGAN & 11.69 & 0.703 & 0.250 & 9.66 & 0.486 & 0.203 & 10.26 & 0.513 & 0.136 \\

    +Prior(R101) & 8.79 & 0.724 & 0.199 & 8.25 & 0.526 & 0.162 & 9.04 & 0.530 & 0.109 \\

    +Prior(InV3)& 7.12 & 0.749 & 0.170 & 7.21 & 0.553 & 0.151 & 8.44 & 0.554 & 0.086 \\
    \shline
  \end{tabular}
  \vspace{0.2cm}
\caption{Comparison of FID, QS, and DDS metrics on three different datasets: FFHQ-256, LSUN-cat, and LSUN-bird.}
\label{table:stylegan_score}
\vspace{-0.2cm}
\end{table}

\begin{table}[t]
\centering
  \begin{tabular}{c|ccc|ccc|ccc
    \hline
    \multirow{2}{*}{ } &
      \multicolumn{3}{c|}{FFHQ-256 } &
      \multicolumn{3}{c|}{LSUN-cat } &
      \multicolumn{3}{c}{LSUN-bird } \\%&

    \cline{2-10}
    &  FID$\downarrow$ &  QS$\uparrow$  & DDS$\downarrow$ & FID$\downarrow$& QS$\uparrow$ & DDS$\downarrow$ &  FID$\downarrow$ &  QS$\uparrow$  & DDS$\downarrow$ \\

    \shline
    StyleGAN & 11.69 & 0.703 & 0.250 & 9.66 & 0.486 & 0.203 & 10.26 & 0.513 & 0.136 \\% & 4.44 & 0.658 & 0.582 \\

    StyleGAN +$\mathcal{L}_Q$ & 9.75 & 0.732 & 0.229 & 8.85 & 0.531 & 0.188 & 9.17 & 0.534 & 0.123 \\

    StyleGAN +RS & 10.00 & 0.715 & 0.196 & 9.11 & 0.505 & 0.149 & 9.79 & 0.521  & 0.095 \\% & 4.33 & 0.660 & 0.584 \\

    \shline
  \end{tabular}
  \vspace{0.14cm}
\caption{Ablation study of our proposed prior with only quality loss and only resampling strategy.}
\label{table:ablation_study}
\vspace{-0.25cm}
\end{table}

\subsection{Prior for Various Datasets}
We also try to apply our proposed prior for various datasets, such as FFHQ-256, LSUN-cat and LSUN-bird. StyleGAN~\cite{karras2019style} recently achieve great success for high-resolution image generation. In this section, we try to combine our method with StyleGAN and evaluate the results on these datasets. We apply two different pre-trained models, ResNet101~\cite{He_2016_CVPR} and InceptionV3~\cite{szegedy2016rethinking} to extract deep features for building the GMM model. We call them +Prior(R101) and +Prior(InV3), respectively. The quantitative results are shown in Table~\ref{table:stylegan_score}. We can observe that our approach can effectively improve the quality(QS), diversity(DDS), and the overall score(FID) over the strong baseline StyleGAN by large margins on all datasets. And the real data prior build on the feature extracted with InceptionV3 model achieves better results than ResNet101, one possible reason is that the evaluation methods are also based on InceptionV3 model. Some visualization results are shown in Figure~\ref{fig:vis_results}.

\begin{table}[h!]
    \begin{minipage}{\columnwidth}
		\begin{minipage}[t]{0.3\columnwidth}
			\small
			\centering
			\begin{tabular}{c|ccc}
			    \hline
				$M$  &  FID$\downarrow$  & QS$\uparrow$ & DDS$\downarrow$ \\
				\shline
				 5 & 7.26 & 0.548 & 0.156 \\

				 7 & 7.21  & 0.553 & 0.151 \\

				 10 & 7.32  & 0.552 & 0.159 \\

				 20 & 7.36  & 0.549 & 0.163 \\

				 50 & 7.45  & 0.543 & 0.177 \\
				 \shline
			\end{tabular}
			\vspace{0.08cm}
			\makeatletter\def\@captype{table}\makeatother\caption{PriorGAN results on LSUN-cat dataset with different number of Gaussian components $M$.}
			\label{table:kernel-densities}
			\vspace{-0.3cm}
		\end{minipage}
		\quad
		\begin{minipage}[t]{0.3\columnwidth}
			\small
			\centering
			\begin{tabular}{c|ccc}
			    \hline
				$\delta$  &  FID$\downarrow$  & QS$\uparrow$ & DDS$\downarrow$ \\
				\shline
				 0.03 & 8.31  & 0.509 & 0.191 \\

				 0.07 & 7.82  & 0.529 & 0.201 \\

				 0.10 & 7.66  & 0.550 & 0.195 \\

				 0.15 & 8.11  & 0.578 & 0.220 \\

				 0.2 & 8.29  & 0.586 & 0.227 \\
				 \shline
			\end{tabular}
			\vspace{0.08cm}
			\makeatletter\def\@captype{table}\makeatother\caption{PriorGAN  results on LSUN-cat dataset with different quality loss weight $\delta$.}
			\label{table:quality-loss-weight}
			\vspace{-0.3cm}
		\end{minipage}
		\quad
		\begin{minipage}[t]{0.3\columnwidth}
			\small
			\centering
			\begin{tabular}{c|ccc}
			    \hline
				$\alpha$  &  FID$\downarrow$  & QS$\uparrow$ & DDS$\downarrow$ \\
				\shline
				 1 & 8.76  & 0.499  & 0.177 \\

				 3 & 8.35 & 0.506 & 0.162 \\

				 5 & 8.30 & 0.514 & 0.143 \\

				 7 & 8.08  & 0.521 & 0.149 \\

				 9 & 8.16 & 0.515 & 0.145 \\
				\shline
			\end{tabular}
			\vspace{0.2cm}
			\makeatletter\def\@captype{table}\makeatother\caption{PriorGAN results on LSUN-cat dataset with different resample weight $\alpha$. }
			\label{table:resample-weight}
			\vspace{-0.3cm}
		\end{minipage}
	\end{minipage}
\end{table}

\subsection{Ablation Study}
To validate the effectiveness of our proposed two key methods: the quality loss and resampling strategy, we apply additional experiments on FFHQ-256, LSUN-cat, LSUN-bird dataset. The quality loss $\mathcal{L}_Q$ and resampling strategy are applied independently to StyleGAN. The resulting models are called StyleGAN +$\mathcal{L}_Q$ and StyleGAN +RS. The results are shown in Table~\ref{table:ablation_study}. We can observe that adding the loss function $\mathcal{L}_Q$ mainly improve QS, which indicates the loss function $\mathcal{L}_Q$ mainly solve the \emph{low quality} problem. Adding the resampling strategy can get a lower DDS, which indicates that the resampling strategy mainly solves the \emph{missing modes} problem. This further validates the effectiveness of our proposed methods.

\subsection{Hyper Parameter Analysis}
In this section, we conduct experiments to investigate the sensitiveness of the number of Gaussian components $M$, quality loss weight $\delta$, and resample weight $\alpha$ in our proposed method. All the experiments are applied to the LSUN-cat dataset.

To reduce the computation cost, we apply PCA on the feature extracted from InceptionV3 model and keep 98\% of the total variance by default. First, we explore how the number of Gaussian components $M$ affects the results. We set $M$ to 5,7,10,20,50 and test the results on LSUN-cat dataset, we show the results in Table~\ref{table:kernel-densities}. The results are not very sensitive to $M$ and we achieve the best performance when setting it to 7. Besides, we set $M$ to 7 and explore how quality loss weight $\delta$ and resample weight $\alpha$ affects the performance, the results are shown in Table ~\ref{table:quality-loss-weight},~\ref{table:resample-weight}.

\section{Conclusion}
In this paper, we analyze current GANs suffer from low quality and missing modes problem. To tackle this problem, we propose a novel method which builds a GMM prior on deep feature space, by using this GMM prior, we propose quality loss to increase generated image quality and resampling strategy to enrich the diversity. Experiments have shown our method can easily insert into any GAN baseline and outperforms state-of-the-art by large margins.

\section*{Broader Impact}
Our PriorGAN is a general method for image generation. From the positive aspect, many generation tasks like image-to-image translation, image editing will benefit from it. We hope our work will serve as a solid baseline and can ease future research in generative models. From the negative aspect, our method can also be applied to some image manipulation techniques, which may cause severe trust issues and security concerns in our society.

{\small
\bibliographystyle{ieee_fullname}