Rewriting a Deep Generative Model: An Overview

"Deep network training is a blind optimization procedure where programmers define objectives but not the solutions that emerge. In this paper, we ask if deep networks can be created in a different way. Can the rules in a network be directly rewritten?"

The Paper | The Code | Interactive Google Colab →



Table of Contents

Why Is Rewriting Deep Generative Model Useful?

We can build a new model without retraining, which is a more involved task.
From the perspective of demystifying deep neural nets, the approach to edit a model gives new insight about the model and how semantic features are captured.
It can also provide some insight into the generalization of deep models to unseen scenarios.
Unlike conventional image editing tools where the desired change is applied on a single image, by editing a GAN, one can apply the edit on every image generated.
Using this tool, one can build new generative models without domain expertise, training time, and computational expense.



An Overview of the Paper

"..we show how to generalize the idea of a linear associative memory to a nonlinear convolutional layer of a deep generator. Each layer stores latent rules as a set of key-value relationships over hidden features. Our constrained optimization aims to add or edit one specific rule within the associative memory while preserving the existing semantic relationships in the model as much as possible. We achieve it by directly measuring and manipulating the model's internal structure, without requiring any new training data."

Try out the interface in Google Colab $\rightarrow$

1. Change Rule With Minimal Collateral Damage

We start with a pre-trained GAN. The authors have used StyleGAN and Progressive GAN pre-trained models trained on multiple datasets.
With a given pre-trained generator(yes we are discarding discriminator), $G(z; θ_0)$, we can generate many(infinite) synthetic images. To generate an image $x_i$ a latent code(just a random vector sampled from a multivariate normal distribution) $z_i$ is required. Thus for a given $z_i$; $x_i = G(z_i; θ_0)$.
The user wants to apply a change such that the new image is $x_{*i}$. For the generator, $G(z_i; θ_0)$ to produce $x_{*i}$ would not be possible since the changed image represents a target distribution that the GAN was not trained on. We thus need to find updated weights $θ_1$ such that $x_{*i} \approx G(z_i; θ_1)$. Interesting!
Note that $θ$ represents the trainable weights or parameters in our GAN generator. The number of parameters is huge in standard SOTA GANs leading to easy overfitting. The aim here is to have a generalized rewritten GAN.
The authors provided two modifications to the standard approach to tackle this manipulation of hidden features in the generator.
Instead of updating all of $θ$, the authors proposed to modify the weights of one layer. Thus all other layers are frozen/made non-trainable.
The objective function(say $L_1$ loss) is applied in the output feature space of that layer instead of the generator's output.
So now, given a layer $L$, let $k$ be the feature output from the $L-1^{th}$ layer(frozen). We can write the feature output from the $L^{th}$ layer as $v = f(k; W_0)$. Thus the output from the target layer, $v$, is defined as a function $f$ with input as $k$ and pre-trained weights $W_0$.
For a latent code $z_i$, the output from the first $L-1$ layer is $k_{*i}$. (Note here the use of $*i$ with $k$ does not mean a change in the rule specified by the user). Thus the output of the target layer $L$ would be $v_i = f(k_{*i}, W_0)$.
For each target example $x_{*i}$, this is specified by the user manually, there is a feature change $v_{*i}$. The aim is to have a generator $G$ such that it can produce target example $x_{*i}$ with minimum interference with other behavior of the generator.
This can be solved by minimizing a simple joint objective function.






Thus, we are updating the weights $W_0$, of target layer $L$ by minimizing smooth and constraint loss. The smooth loss function ensures that the output generated by the new rule is not far apart from the initial(pre-trained) one. The constraint loss ensues a specific rule to be modified or added.
Also $|| . ||^2$ denotes the $L2$ loss.

2. The Analogy of Associative Memory

Any matrix $W$ can be used as an associative memory that stores a set of key-value pairs {$(k_i, v_i)$}. This stored set of key-value pair can be retrieved by matrix multiplication such that,
$v_i \approx Wk_i$
The reason for approximate quality is because practically such a memory bank is not error-free. We can create error-free memory by having keys such that they form a set of mutually orthogonal unit-norm vectors. However, in the above formulation of associative memory, only $N$ orthogonal keys.
The authors have considered the convolutional weights of the target layer $L$ as an associative memory. Thus instead of thinking the layer as a collection of convolutional filtering operations, **the layer is considered as a memory that associates memory keys to values.
Each key is a single location feature vector. At the same time, the value is the output arrangement of pixels.
To support more than $N$ nonorthogonal keys, {$k_i$}, instead of requiring exact equality that is $v_i = Wk_i$, error be minimized such that,
$W_0 \overset{\Delta}{=} \underset{x}{\arg\min} \sum_{i} ||v_i - Wk_i||^2$

This is a least-squares problem, the minimal solution of which can be found by solving for $W_0$ using the normal equation $W_0KK^T = VK^T$. Here $K$ and $V$ are simplified notations whose i-th column is the i-th key or value. More on least-square approximation here.

3. How to Update $W$ to Insert a new Value?

It should store a new value.
It should continue to minimize errors in all the previously stored values.

For the requested mapping $k_* \rightarrow v_*$ the final form of the equation transforms the soft error minimization objective into hard constraint such that the weights be updated in a particularly straight direction $C^{-1}k_*$.
The update direction $C^{-1}k_*$ is determined only by overall key statistics and the specific targeted key $k_*$.

4. Generalizing to a Non-Linear layer

They first define the update direction, $d \overset{\Delta}{=} C^{-1}k_*$. It was derived in the previous section.
Then, in order to insert a new key $k_* \rightarrow v_*$, they begin with $W_0$ and perform an optimization over the rank-one subspace defined by the row vector $d^T$. First, this optimization is solved:
$\Lambda_1 = \underset{\Lambda\in\R^{M}}{\arg\min} ||v_* - f(k_*; W_0 + \Lambda d^T)||$
Once $\Lambda_1$ is calculated, weight is updated such that,
$W_1 = W_0 + \Lambda_1 d^T$

User Interface

Figure 3: The copy-paste-context interface for rewriting a model. (Source)

Copy - the user can use a brush to copy fascinating region in an image to define target value $V_*$
Paste - the user can paste the copied value into a single target image specifying $K_* \rightarrow V_*$ constraint.
Context - for generalization, the user can select the same interesting region in multiple images. This establishes the updated direction $d$ for the associative memory.

Try out the interface in Google Colab $\rightarrow$

Figure 4: User interface as seen after running the linked Colab notebook.

Putting it all Together and Ending Note

Putting objects/interesting regions into a new context. You can add trees to the top of the building or a dome in place of a pointy tower. You can also put the moustache in place of an eyebrow!
Removing undesired features like a watermark in an image or some artifacts in an image. This is a handy outcome of this idea.

"Machine learning requires data, so how can we create effective models for data that do not yet exist? Thanks to the rich internal structure of recent GANs, in this paper, we have found it feasible to create such models by rewriting the rules within existing networks."