TLDR; The authors augment the A3C (Asynchronous Actor Critic) algorithm with auxiliary tasks. These tasks share some of the network parameters but value functions for them are learned off-policy using n-step Q-Learning. The auxiliary tasks only used to learn a better representation and don't directly influence the main policy control. The technique, called UNREAL (Unsupervised Reinforcement and Auxiliary Learning), outperforms A3C on both the Atari and Labyrinth domains in terms of performance and training efficiency.


#### Key Points

- Environments contain a wide variety of possible training signals, not just cumulative reward
- Base A3C agent uses CNN + RNN
- Auxiliary Control and Prediction tasks share the convolutional and LSTM network for the "base agent". This forces the agent to balance improvement and base and aux. tasks.
- Auxiliary Tasks
  - Use off-policy RL algorithms (e.g. n-step Q-Learning) so that the same stream of experience from the base agent can be used for maximizing all tasks. Experience is sampled from a replay buffer.
  - Pixel Changes (Auxiliary Control): Learn a policy for maximally changing the pixels in a grid of cells overlaid over the images
  - Network Features (Auxiliary Control): Learn a policy for maximally activating units in a specific hidden layer
  - Reward Prediction (Auxiliary Reward): Predict the next reward given some historical context. Crucially, because rewards tend to be sparse, histories are sampled in a skewed manner from the replay buffer so that P(r!=0) = 0.5. Convolutional features are shared with the base agent.
  - Value Function Replay: Value function regression for the base agent with varying window for n-step returns.
- UNREAL
  - Base agent is optimized on-policy (A3C) and aux. tasks are optimized off-policy.
- Experiments
  - Agent is trained with 20-step returns and aux. tasks are performed every 20 steps.
  - Replay buffer stores the most recent 2k observations, actions and rewards
  - UNREAL tends to be more robust to hyperparameter settings than A3C
  - Labyrinth
    - 38% -> 83% human-normalized score. Each aux. tasks independently adds to the performance.
    - Significantly faster learning, 11x across all levels
    - Compared to input reconstruction technique: Input reconstruction hurts final performance b/c it puts too much focus on reconstructing relevant parts.
  - Atari
    - Not all experiments are completed yet, but UNREAL already surpasses state of the art agents and is more robust.

#### Thoughts

- I want an algorithm box please :)

Single cells behave in complex and sometimes seemingly random ways. We have applied Generative Adversarial Networks (GANs), a form of artificial intelligence, to make sense of the way genes are controlled in skin cells. A GAN involves two separate neural networks, a ‘generator’ and a ‘discriminator’. The generator simulates cells and the discriminator tries to tell the difference between the fake data created by the simulator, and data from real cells. As they compete against each other they improve at their tasks and provide new insights into the way cells behave.   

Our neural network approach allows us to understand the relationship between different genes and how this contributes to cell behaviour. One of the networks, the generator, is responsible for simulating cells and we use this to predict how genes are controlled under different conditions, effectively simulating what would previously have been laborious and painstaking experiments. GANs also make it possible to compare data from multiple labs produced in different conditions, opening up the opportunity to answer new questions.

This paper implements a parameter expansion technique with Markov Chain Monte Carlo (MCMC) for the multivariate probit model (MVP). Since the covariance matrix is not identifiable from multivariate ordinal data, a correlation matrix restriction is appropriate with MVP. This is a problem for Bayesian inference, as the structural constraints on correlation matrices make the selection of a good prior and the estimation of a posterior difficult.   

The main advantage of parameter expansion with this model is that we can draw the covariance matrix in the expanded parameter space from an inverse Wishart distribution with a simple Gibbs sampler. At each iteration, all parameters are drawn from the expanded space and rescaled to their original form. Since the parameter draws are only dependent on the latest latent variable, the dependence between iterations is reduced. This technique improves the mixing and convergence rate of the MCMC.

Parameter expansion for the MVP model is also easy to implement for stochastic search variable selection (SSVS). An indicator variable $\lambda_i$ is included as a parameter, where each 1 in $\lambda_i$ is associated to a covariate selected for the model.

Gibbs sampler algorithm:

Let $q$-variate vector $Z_i \sim MVN( \beta X_i, R )$ be a latent variable, truncated in all dimensions according to the corresponding cutpoints, where $R$ is a $q \times q$ correlation matrix.

Let $\beta$ be a $q \times p$ regression coefficient matrix, $\gamma$ be the set of cutpoints where $\gamma_0 = -\infty$, $\gamma_1 = 0$ and $\gamma_k = \infty$, and $k$ be the number of categories for the response.


Map $Z $ onto $W$:

Cycle through draws of the univariate truncated $Z_{i,j} \sim N(\mu_{i,j} - \frac{1}{R^{-1}_{j,j}}R^{-1}_{j,-j}(Z_{i,-j}-\mu_{i,-j}),  \frac{1}{R^{-1}_{j,j}})$, where $\mu_i = \beta X_i$, to get the truncated multivariate normal distribution of $W_i$.

$R \rightarrow \Sigma$:

Draw $\Sigma | Z \sim IW(n-p+q-1, WW^T - WX^T(XX^T)^{-1}XW^T)$

$\beta \rightarrow \alpha$:

Draw $\alpha | W, \Sigma \sim N_{q,p}(WX^T(XX^T)^{-1}, \Sigma \otimes (XX^T)^{-1})$

where $N_{q,p}$ is a matrix variate normal distribution.

$\gamma \rightarrow \theta$:

Draw $\theta_{j,k+1} | W, Y \sim Unif(max_i(\{W_{ij}|Y_{ij}=c\}), min_i(\{W_{ij}|Y_{ij} = c + 1\})) $ for all free cutpoints.

Rescale the parameters back to the original space:

$Z_{i,j} = \dfrac{W_{i,j}}{\sqrt{\Sigma_{j,j}}}$

$R_{i,j} = \dfrac{\Sigma_{i,j}}{\sqrt{\Sigma_{i,i}\cdot\Sigma_{j,j}}}$

$\beta_{k,j} = \dfrac{\alpha_{k,j}}{\sqrt{\Sigma_{j,j}}}$

$\gamma_{j,k+1} = \dfrac{\theta_{j,k+1}}{\sqrt{\Sigma_{j,j}}}$

Repeat.

There are mathematicians, still today, who look at deep learning, and get real salty over the lack of convex optimization. That is to say: convex functions are ones where you have an actual guarantees that gradient descent will converge,  and mathematicians of olden times (i.e. 2006) spent reams of paper arguing that this or that function had convex properties, and thus could be guaranteed to converge, under this or that set of arcane conditions. And then, Deep Learning came along, with its huge, nonlinear, very much nonconvex objective functions, that it was nonetheless trying to optimize via gradient descent. From the perspective of an optimization theorist, this had the whiff of heresy, but exceptionally effective heresy. And, so, the field of DL has half-exploded, half-stumbled along, showcasing a portfolio of very impressive achievements, but with theory very much a secondary priority relative to performance. 

Something else that gradient descent isn’t supposed to be able to do is learn models that include discrete (i.e. non-continuous) operators. Without continuous gradients, the functions don’t have an obvious way to “push” in a certain direction, to modulate the loss at the end of the network. Discrete nodes mean that the value just jumps from being in one state, to being in the other, with no intermediate values. This has historically posed a problem for algorithms fueled by gradient descent.


The authors of this paper came up with a solution that is 60% cleverness, and 40% just guessing that “even if we ignore the theory, things will probably work well enough”. But, first, their overall goal: to create a Variational Auto Encoder where the latent states, the compressed internal representation that is typically an array of continuous values, is instead an array of categorical values. The goal of this was 1) to have a representation type that was a better match for the discrete nature of data types like speech (which has distinct phonemes we might like to discretely capture), and, 2) to have a more compressed latent space that would (of necessity) focus on more global information, and leave local pixel-level information to be learned by the expressive PixelCNN decoder. The way they do this is remarkably simple. First, they learn a typical VAE encoder, mapping from the input pixels to a continuous z space. (An interesting sidenote here is that this paper uses spatially organized z; instead of using one single z vector to represent the whole image, they may have 32x32 spatial locations, each of which has its own z vector,  to represent at 128x128 image). Then, for each of the spatial regions, they take the continuous vector produced by the network, and compare it to a fixed set of “embedding” vectors, of the same shape. That spatial location is then lumped into the category of the embedding that it’s closest to, meaning that you end up with a compressed layer of 32x32 (in this case) spatial regions, each of which is represented by a categorical number between 0 and max-num-categories. Then, the network passes forward the embedding that this input vector was just “snapped” to being, Then, the decoder uses the full spatial location set of embeddings to do its decoding.
 
https://i.imgur.com/P8LQRYJ.png

The clever thing here comes when you ask how to train the encoder to produce a different embedding, when there was this discrete “jump” that happened. The authors choose to just avoid the problem, more or less. They do that by just taking the gradient signals that come back from the end of the network to the embedding, and just pass those directly to the vector that was used to nearest-neighbors-lookup the embedding. Basically, they pretend that they passed the vector through the rest of the network, rather than the embedding. The embeddings are then trained in a K Means Clustering kind of way; with the embeddings being iteratively updated to be closer to the points that were assigned to their embedding in each round of training. This is the “Vector Quantization” part of VQ-VAE

Overall, this seems to perform quite well: with the low capacity of the latente space meaning that it is incentivized to handle more global structure, while leaving low level pixel details to the decoder. It is also much easier to fit after-the-fact distributions over; once we’ve trained a VQ-VAE, we can easily learn a global model that represents the location by location dependencies between the categories (i.e. a 1 in this corner means at 5 in this other corner is more probable). This gives us the ability to have an analytically specified distribution, in latent space, that actually represents the structure of how these “concept level categories” relate to each other. By contrast, with most continuous latent spaces, it’s intractable to learn an explicit density function after the fact, and thus if we want to be able to sample we need to specify and enforce a prior distribution over z ahead of time. 

Unsupervised representation learning is a funny thing: our aspiration in learning representations from data is typically that they’ll be useful for future tasks, but, since we (by definition) don’t have access to labels, our approach has historically been to define heuristics, such as representing the data distribution in a low-dimensional space, and hope that those heuristics translate to useful learned representations. And, to a fair extent, they have. However, this paper’s goal is to attach this problem more directly, by explicitly meta-learning an unsupervised update rule so that performs well in future tasks. They do this by: 

https://i.imgur.com/EEkpW9g.png

1) Defining a parametrized weight update function, to slot into the role that Stochastic Gradient Descent on a label-defined loss function would play in a supervised network. This function calculates a “hidden state”, is defined for each neuron in each layer, and takes in the pre and post-nonlinearity activations for that batch, the hidden state of the next layer, and a set of learned per-layer “backwards weights”. The weight update for that neuron is then calculated using the current hidden state, the last batch's hidden state, and the current value of the weight. In the traditional way of people in this field who want to define some generic function, they instantiate these functions as a MLP. 

2) Using that update rule on the data from a new task, taking the representing resulting from applying the update rule, and using it in a linear regression with a small number of samples. The generalization performance from this k-shot regression, taken in expectation over multiple tasks, acts as our meta training objective. By back-propagating from this objective, to the weight values of the representation, and from there to the parameters of the optimization step, they incentivize their updater to learn representations that are useful across some distribution of tasks. 


A slightly weird thing about this paper is that they train on image datasets, but shuffle the pixels and use a fully connected network rather than a conv net. I presume this has to do with the complexities of defining a weight update rule for a convolution, but it does make it harder to meaningfully compare with other image-based unsupervised systems, which are typically done using convolution. An interesting thing they note is that, early in meta-training on images, their update rules generalize fairly well to text data. However, later in training the update rules seem to have specialized to images, and generalize more poorly to images.