Introduction

I originally wanted to write down the proof for the Gumbel-max trick but soon realized it is actually the same idea as a much more common problem: exponential race. So, in this note let’s go from this common problem and arrive at the Gumbel-max trick.

Competing Alarms

As a preparation let’s solve a probability problem first.

There are \(N\) clocks started simultaneously, such that the \(i\)-th clock rings after a random time \(T_i \sim \text{Exp}(\lambda_i)\)

• (1) Designate \(X\) as the random time after which some clock (i.e any one of the clocks) rings, find the distribution of \(X\)

• (2) Find the probability of the \(i\)-th clock rings first

Let \(X = \min \{T_1, T_2, \dots, T_n \}\) and \(F_X(t)\) and \(F_{T_i}(t)\) be the CDFs. We also have \(F_{T_i}(t) = 1- e^{-\lambda_it}\).

Following order statistics of \(\min\), we have \(P(X>t) = \prod_{i=1}^N P(T_i>t)\) or equivalently,

Therefore

i.e. the \(\min\) of a set of i.i.d exponential random variables is still an exponential random variable with the rate \(\lambda_X\) being the sum of the rates of that set of random variables.

For the second part of the problem, we can consider two competing alarms \(T_1\) and \(T_2\) to begin with. Our goal is to find \(P(T_1<T_2)\).

Now, let’s consider one specific clock \(T_k\) versus all the rest, noted as \(T_{-k} = \min \{T_i\}_{i \neq k}\). According to \(\ref{part1}\), we know that \(T_{-k} \sim \text{Exp}(\sum_{i\neq k} \lambda_i)\). Using the result above we have the solution for part (2) as follows

Of course, we can do the integration directly and get the same result

By the way, this setup with multiple exponential random variables and we look for the first arrival is also called exponential race.

Exponential-Min Trick

I just made up the name “Exponential-Min”. The better name for this section is probably Sampling from Multinomial by Argmining.

Suppose we have a set of positive numbers \([\lambda_1, \lambda_2, \lambda_3, \dots, \lambda_N]\). Correspondingly we have a normalized probabiilty vector \(\vec{p}=[p_1, p_2, p_3, \dots, p_N]\), where \(p_k = \frac{\lambda_k}{\sum_{i=1}^N\lambda_i}\). This probability vector specifies a multinominal distribution over \(N\) choices.

Now, if we were to get a sample \(\{1, 2, \dots, N\}\) according to this multinominal distribution specified by \(\vec{p}\) (which is fundamentally specified by \([\lambda_1, \lambda_2, \lambda_3, \dots, \lambda_N]\)), what should we do?

Normally, we do the following:

1. We have \([\lambda_1, \lambda_2, \lambda_3, \dots, \lambda_N]\)
2. We compute \(\vec{p}=[p_1, p_2, p_3, \dots, p_N]\), where \(p_k = \frac{\lambda_k}{\sum_{i=1}^N\lambda_i}\).
3. We generate a uniform random number \(Q\) between 0 and 1, i.e. \(Q \sim \text{Uniform}(0,1)\)
4. We figure out where \(Q\) lands, i.e. if \(p_i < Q < p_{i+1}\) we return \(i\). (Of couse we should use \(p_0=0\) and \(p_{N+1}=1\))

But that’s the boring way. Now we have this new Exponential-Min trick, we can do the following:

1. We have \([\lambda_1, \lambda_2, \lambda_3, \dots, \lambda_N]\)
2. We don’t compute \(\vec{p}\); instead we sample \(T_i \sim \text{Exp}(\lambda_i)\) for \(i=1, 2, \dots, N\), i.e. we have a total of \(N\) samples, one from each \(\text{Exp}(\lambda_i)\)
3. We now take \(\argmin([T_1, T_2, \dots, T_N])\) as our result sample
4. We proved in \(\ref{part2}\) that such a result sample indeed follows multinominal distribution specified by \(\vec{p}=[p_1, p_2, p_3, \dots, p_N]\), where \(p_k = \frac{\lambda_k}{\sum_{i=1}^N\lambda_i}\).

Thus, somehow we ended up Sampling from Multinomial by Argmining!

Gumbel Distribution

Now let’s move on to Gumbel distribution from Exponential distribution.

Gumbel distribution with unit scale (\(\beta=1\)) is parameterized by location parameter \(\mu\). \(\text{Gumble}(\mu)\) has CDF and PDF as follows

Given a set of \(N\) independent Gumbel random variables \(G_i\), each with their own parameter \(\mu_i\), i.e. \(G_i \sim \text{Gumbel}(\mu_i)\).

Gumbel distribution has two properties that are quite analogous the exponential race example above.

• (1) Let \(Z = \max \{G_i \}\), then \(Z \sim \text{Gumble}\left(\mu_Z = \log \sum_{i=1}^N e^{\mu_i} \right)\)

The proof is straightforward and similar to above:

• (2) A corollary of the above is that the probability of \(Z_k\) being the max is \(P(Z_k > Z_{-k}) = \frac{e^{\mu_k}}{\sum_{i=1}^N e^{\mu_i}}\)

Gumbel-Max Trick

Now here we can tell nearly an identical/parallel story as in the section Exponential-Min Trick. And, this section should really be called Sampling from Multinomial by Argmaxing.

The main differences are

• The numbers (parameters) \(\mu_i\) can be potentially negative, whereas \(\lambda_i\) must be positive
• The probability vector is determined by \(p_k = \frac{e^{\mu_k}}{\sum_{i=1}^N e^{\mu_i}}\) instead of \(p_k = \frac{\lambda_k}{\sum_{i=1}^N\lambda_i}\)
• We generate samples with \(G_i \sim \text{Gumbel}(\mu_i)\) instead of \(T_i \sim \text{Exp}(\lambda_i)\)
• We take \(\argmax\) over \(G_i\) instead of taking \(\argmin\) over \(T_i\)

When is Gumbel-Max Trick Useful?

It seems a lot of work to sample multinominal by argmaxing over Gumble samples (or argmining over Exponential samples). In what situation do we ever want to do this?

The short answer is that Gumble-Max trick allows us to make a sampling step differentiable. Specifically, it makes sampling from multinomial distribution differentiable. We’ll take a closer look at this in a future post but pause for a second and think about it. We are saying it is possible to differentiate through the action of drawing a discrete sample from a multinomial distribution! This was a pretty surprising/amazing possibility to me.

Regarding downstream applications, differentiating through sampling is an important “trick” in neural network based variational inference in general. Multinomial discrete random variables are prevalent in many learning problems. Gumble-max trick allows us to work with them in many interesting neural variational inference problems, which we will look into in future posts.

Overview

This page is a high-level summary / notes of various recent results in language modeling with little explanations. Papers to cover are as follows:

[1] AWD Language Model

• Merity, Stephen, Nitish Shirish Keskar, and Richard Socher. “Regularizing and optimizing LSTM language models.” arXiv preprint arXiv:1708.02182 (2017).

[2] Neural Cache

• Grave, Edouard, Armand Joulin, and Nicolas Usunier. “Improving neural language models with a continuous cache.” arXiv preprint arXiv:1612.04426 (2016).

[3] Dynamic Evaluation

• Krause, Ben, et al. “Dynamic evaluation of neural sequence models.” arXiv preprint arXiv:1709.07432 (2017).

[4] Memory-based Parameter Adaptation (MbPA)

• Sprechmann, Pablo, et al. “Memory-based parameter adaptation.” arXiv preprint arXiv:1802.10542 (2018).

[5] Hebbian Softmax

• Rae, Jack W., et al. “Fast Parametric Learning with Activation Memorization.” arXiv preprint arXiv:1803.10049 (2018).

[6] Higher-rank LM / Mixture-of-Softmax (MoS)

• Yang, Zhilin, et al. “Breaking the softmax bottleneck: A high-rank RNN language model.” arXiv preprint arXiv:1711.03953 (2017).

This is by no means an exhaustive literature review - they are only a selection of a few of the most recent state-of-the-art results. AWD LM [1] has almost become the de-facto baseline LM for many of the other papers, where the main innovations area special version of Averaged SGD (ASGD) along with DropConnection based Weight Dropping regularization in the hidden -to-hidden mapping of a LSTM model.

It has been found a global LM is ineffective in reacting to local patterns at test time, such as once a rare word appears furthe reappearance in the peoximty is much more likely than predicted by a global LM. To allow for faster reaction to local patterns, [2 - 5] propose various schemes involving a fast-learning non-parametric component and blend its predictions or parameters with the global learned parametric LM. A quick comparsion of these 4 papers are in the table below.

Table 1. Comparison of methods in Ref [2-5]

And finally, [6] highlights and mostly solved a fairly general problem of softmax over product produced by rank-limited matrices which is common in the decoder in a LM.

AWD LM

Neural Cache

Dynamic Evaluation

Memory-based Parameter Adaptation (MbPA)

Hebbian Softmax

Higher-rank LM / Mixture-of-Softmax (MoS)

A short list of resources and topics covering the essential quantitative tools for Data Scientists, Machine Learning Engineers/Scientists, Quant Developers/Researchers and those who are preparing to interview for these roles.

At a high-level we can divide things into 3 main areas:

1. Machine Learning
2. Coding
3. Math (calculus, linear algebra, probability, etc)

Depending on the type of roles, the emphasis can be quite different. For example, AI/ML interviews might go deeper into the latest deep learning models, while quant interviews might cast a wide net on various kinds of math puzzles. Interviews for research-oriented roles might be lighter on coding problems or at least emphasize on algorithms instead of software designs or tooling.

List of resources

A minimalist list of the best/most practical ones:

Machine Learning:

• Course on classic ML: Andrew Ng’s CS229 (there are several different versions, the Cousera one is easily accessible. There is also an older version recorded at Stanford)
• Book on classic ML: Alpaydin’s Intro to ML link
• Course with a deep learing focus: CS231 from Stanford, lectures available on Youtube.

If you are just breaking into the field I think the above are enough, stop there and move on to other areas of preparation. Here are a few very optional items, mostly on deep learning, in case you have more time:

• Overview book on deep learning: Neural Networks and Deep Learning by Michael Nielson.
• Amazing book on deep laerning for NLP: Neural Network Methods for Natural Language Processing by Yoav Goldberg
• Pick one of those Udacity nanodegrees on deep learning / self-driving cars to get some hands on experience with deep learning frameworks (Tensorflow, Pytorch, MXNet)

Coding:

• Course: MIT OCW 6006 link
• Book: Cracking the Coding Interview link
• Practice sites: Leetcode, HackerRank
• SQL tutorial: from Mode Analytics

Math:

• Calculus and Linear Algebra: undergrad class would be the best, refresher notes from CS229 link
• Probability: Harvard Stats110 link; book from the same professor.
• Statistics: Shaum’s Outline link.
• [Optional] Numerical Methods and Optimization: these are two different topics really, college courses are probably the best bet. I have yet to find good online courses for them. But don’t worry, most interviews won’t really touch on them.

List of topics

Here is a list of topics from which interview questions are often derived. The depth and trickiness of the questions certainly depend on the role and the company.

Under topic I try to add a few bullet points of the key things you should know.

Machine learning

• Models (roughly in decreasing order of frequency)
  • Linear regression - e.g. assumptions, multicollinearity, derive from scratch in linear algebra form
  • Logistic regression
    • be able to write out everything from scratch: from definitng a classficiation problem to the gradient updates
  • Decision trees/forest - e.g. how does a tree/forest grow, on a pseudocode level
  • Clustering algorithms
    • e.g. K-means, agglomerative clustering
  • SVM
    • e.g. margin-based loss objectives, how do we use support vectors, prime-dual problem
  • Generative vs discriminative models
    • e.g. Gaussian mixture, Naive Bayes
  • Anomaly/outlier detection algorithms (DBSCAN, LOF etc)
  • Matrix factorization based models
• Training methods
  • Gradient descent, SGD and other popular variants - Understand momentum, how they work, and what are the diffrences between the popular ones (RMSProp, Adgrad, Adadelta, Adam etc) - Bonus point: when to not use momentum?
  • EM algorithm - Andrew’s lecture notes are great, also see this
  • Gradient boosting
• Learning theory / best practice (see Andrew’s advice slides)
  • Bias vs variance, regularization
  • Feature selection
  • Model validation
  • Model metrics
  • Ensemble method, boosting, bagging, bootstraping
• Generic topics on deep learning
  • Feedforward networks
  • Backpropagation and computation graph
    • I really liked the miniflow project Udacity developed
    • In addition, be absolutely familiar with doing derivatives with matrix and vectors, see Vector, Matrix, and Tensor Derivatives by Erik Learned-Miller and Backpropagation for a Linear Layer by Justin Johnson
  • CNN, RNN/LSTM/GRU
  • Regularization in NN, dropout, batch normalization

Coding essentials

There are a lot of resources online on how to prepare for this, some are already listed in the resources section. I think the key thing is to pick a language and know it really well. For example, for Python, if you claim you know it well and have used in some non-trivial code base. I’ll assume you know why abstract base class exists, how decorator works in general and what @property means along with a few language-specific data structures (like OrderedDict, deque, defaultdict and etc).

The bare minimum of coding concepts you need to know well.

• Data structures:
  • array, dict, link list, tree, heap, graph, ways of representing sparse matrices
• Sorting algorithms:
  • see this from brilliant.org
  • I think in real life you’ll most likely never implement a sorting algorithm, but I think quick sort is very cool and quick select / partition is used in many other places, so take a look.
• Tree/Graph related algorithms
  • Traversal (BFS, DFS)
  • Shortest path (two sided BFS, dijkstra)
• Trees related
  • TBA
• Heap related
  • TBA
• Recursion and dynamic programming
  • TBA

Calculus

Just to spell things out

• Derivatives
  • Product rule, chain rule, power rule, L’Hospital’s rule,
  • Partial and total derivative
  • Things worth remembering
    • common function’s derivatives
    • limits and approximations
  • Applications of derivatives: e.g. this
• Integration
  • Power rule, integration by sub, integration by part
  • Change of coordinates
• Taylor expansion
  • Single and multiple variables
  • Taylor/McLauren series for common functions
  • Derive Newton-Raphson
• ODEs, PDEs (common ways to solve them analytically)

Linear algebra

• Vector and matrix multiplication
• Matrix operations (transpose, determinant, inverse etc)
• Types of matrices (symmetric, Hermition, orthogonal etc) and their properties
• Eigenvalue and eigenvectors
• Matrix calculus (gradients, hessian etc)
• Useful theorems
• Matrix decomposition
• Concrete applications in ML and optimization

Probability

Solving probability interview questions is really all about pattern recognition. To do well, do plenty of exercise from this and this. This topic is particularly heavy in quant interviews and usually quite light in ML/AI/DS interviews.

• Basic concepts
  • Event, outcome, random variable, probability and probability distributions
• Combinatorics
  • Permutation
  • Combinations
  • Inclusion-exclusion
• Conditional probability
  • Bayes rule
  • Law of total probability
• Probability Distributions
  • Expectation and variance equations
  • Discrete probability and stories
  • Continuous probability: uniform, gaussian, poisson
• Expectations, variance, and covariance
  • Linearity of expectation
    • solving problems with this theorem and symmetry
  • Law of total expectation
  • Covariance and correlation
  • Independence implies zero correlation
  • Hash collision probability
• Universality of Uniform distribution
  • Proof
  • Circle problem
• Order statistics
  • Expectation of min and max and random variable
• Graph-based solutions involving multiple random variables
  • e.g. breaking sticks, meeting at the train station, frog jump (simplex)
• Approximation method: Central Limit Theorem
  • Definition, examples (unfair coins, Monte Carlo integration)
  • Example question
• Approximation method: Poisson Paradigm
  • Definition, examples (duplicated draw, near birthday problem)
• Poisson count/time duality
  • Poisson from poissons
• Markov chain tricks
  • Various games, introduction of martingale

Statistics

• Z-score, p-value
• t-test, F-test, Chi2 test (know when to use which)
• Sampling methods
• AIC, BIC

[Optional] Numerical methods and optimization

• Computer errors (e.g. float)
• Basic root finding (newton method, bisection, secant etc)
• Interpolating
• Numerical integration and difference
• Numerical linear algebra
  • Solving linear equations, direct methods (understand complexities here) and iterative methods (e.g. conjugate gradient), maybe BFGS
  • Matrix decompositions/transformations (e.g. QR, Givens, LU, SVD etc)
  • Eigenvalue solvers (e.g. power iteration, Arnoldi/Lanczos etc)
• ODE solvers (explicit, implicit)
• Finite-difference method, finite-element method
• Optimization topics: linear programming (and convex opt in general), calculus of variations

Overview

This page is a high-level summary of various generative models with little explanations. Models to cover are as follows:

Variational Autoencoders (VAE)

• Kingma, Diederik P., and Max Welling. “Auto-encoding variational bayes.” arXiv preprint arXiv:1312.6114 (2013).

Adversarial Variational Bayes (AVB)

Extention to VAE to use non-Gaussian encoders

• Mescheder, Lars, Sebastian Nowozin, and Andreas Geiger. “Adversarial variational bayes: Unifying variational autoencoders and generative adversarial networks.” arXiv preprint arXiv:1701.04722 (2017).

Generative Adverserial Networks (GAN)

• Goodfellow, Ian, et al. “Generative adversarial nets.” Advances in neural information processing systems. 2014.

Generalized divergence minimization GAN (\(f\)-GAN)

• Nowozin, Sebastian, Botond Cseke, and Ryota Tomioka. “f-gan: Training generative neural samplers using variational divergence minimization.” Advances in Neural Information Processing Systems. 2016.

Wasserstein GAN (WGAN)

• Arjovsky, Martin, Soumith Chintala, and Léon Bottou. “Wasserstein gan.” arXiv preprint arXiv:1701.07875 (2017).

Adversarial Autoencoders (AAE)

• Makhzani, Alireza, et al. “Adversarial autoencoders.” arXiv preprint arXiv:1511.05644 (2015).

Wasserstein Auto-Encoder (WAE)

• Tolstikhin, Ilya, et al. “Wasserstein Auto-Encoders.” arXiv preprint arXiv:1711.01558 (2017).

Cramer GAN

• Bellemare, Marc G., et al. “The Cramer Distance as a Solution to Biased Wasserstein Gradients.” arXiv preprint arXiv:1705.10743 (2017).

VAE

Model setup:

• Recognition model: \(q_\phi(z \vert x) = \mathcal N(\mu=h_1(x), \sigma^2 \mathbf I=h_2(x)\mathbf I)\)
• Assumed fixed prior: \(p(z) = \mathcal N(0,\mathbf I)\)
• Generation model: \(p_\theta(x \vert z) = \mathcal N(\mu=g_1(z), \sigma^2 \mathbf I=g_2(z)\mathbf I)\)
  • Implied (but intractable) posterior: \(p_\theta(z\vert x)\)

Key equations:

Optimization objective:

Gradient-friendly Monte Carlo:

Difficulties in calculating \(\mathcal L(\theta, \phi, x^i)\):

• Due to the generality of \(q\) and \(p\) (typically a neural network), the expectation in \(\ref{vae_elbo}\) does not have an analytical form. So we need to resort to Monte Carlo estimation.
• Furthermore, direct sampling \(z\) according to \(q\) poses difficulty in taking derivative against parameters \(\phi\) that parameterizes the distribution \(q\).

Solution: Reparameterization Trick

Find smooth and invertible transformation \(z=g_\phi(\epsilon)\) such that with \(\epsilon\) drawn from a fixed (non-parameterized) distribution \(p(\epsilon)\) we have \(z \sim q(z; \phi)\), so

For the Normal distribution used here (\(q_\phi(z\vert x)\)), it is convenient to use location-scale transformation, \(z=\mu+\sigma * \epsilon\) with \(\epsilon \sim \mathcal N(0,\mathbf I)\).

For total \(N\) data points with mini batch size \(M\):

For sufficiently large batch size \(M\), the inner loop sample size \(L\) can be set to 1. Due to stochastic mini batch gradient descent and stochastic expectation estimation, this is also called doubly stochastic estimation.

Using non-Gaussian encoders

Todo: discuss AVB paper

Gumble trick for discrete latent variables

Ref for this section:

1. Gumble max trick https://hips.seas.harvard.edu/blog/2013/04/06/the-gumbel-max-trick-for-discrete-distributions/
2. Balog, Matej, et al. “Lost Relatives of the Gumbel Trick.” arXiv preprint arXiv:1706.04161 (2017).
3. Jang, Eric, Shixiang Gu, and Ben Poole. “Categorical reparameterization with gumbel-softmax.” arXiv preprint arXiv:1611.01144 (2016).

Gumble distribution:

\(f\)-GAN and GAN

Prelude on \(f\)-divergence and its variational lower bound

The f-divergence family

where the generator function \(f: \mathbb{R}_{+} \rightarrow \mathbb{R}\) is a convex, lower-semicontinuous function satisfying \(f(1) = 0\).

Every convex, lower-semicontinuous function has a convex conjugate function \(f^c\), also known as Fenchel conjugate. This function is defined as

Function \(f^c\) is again convex and lower-semicontinuous and the pair \((f,f^c)\) is dual to each other, i.e. \(\left(f^{c}\right)^c=f\). So we can represent \(f\) as

With this we can establish a lower bound for estimating the f-divergence in general

where \(\mathcal T\) is an arbitrary class of functions \(T: \mathcal X \rightarrow \mathbb R\). The inequality is due to Jessen’s inequality and constraints imposed by \(\mathcal T\).

The bond is tight for

Generative adversarial training

Suppose our goal is to come up with a distribution \(Q\) (model) that is close to \(P\) (the data distribution) and the similarity score (loss) is measured by \(D_f(P \| Q)\). However the direct calculation of \(\ref{f_div}\) is intractable, such as the case where the functional form of \(P\) is unknown and \(Q\) is a complex model parameterized by a neural network.

To be specific:

• Evaluating \(q(x)\) at any \(x\) is easy, but integrating it is hard due to lack of easy functional form.
• For \(p(x)\), we do not know how to evaluate it at any \(x\)
• Sampling from both \(P\) and \(Q\) are easy. Because drawing from data set approximates \(x \sim P\) and we can make the model \(Q\) take random vectors as input which are easy to produce.

Fortunately, we can sample from both of them easily. In this case, \(\ref{f_lowerbound}\) offers a way to estimate the lower bound of the divergence. We would need to maximize this lower bound by changing \(T\) so that it is close to the true divergence, then minimize it over \(Q\). This is formally stated as follows.

To ensure that the output of \(T_\omega\) respects the domain of \({f^c}\), we define \(T_\omega(x) = g_f(V_\omega(x))\), where \(V_\omega: \mathcal X \rightarrow \mathbb R\) without any range constraints on the output and \(g_f: \mathbb R \rightarrow \text{dom}_{f^c}\) is an output activation function specific to the \(f\)-divergence used with suitable output ranges.

GAN

For the original GAN, with a divergence target similar to Jensen-Shannon \(\begin{equation} F(\theta, \omega) = \mathbb{E}_{x\sim P}[\log D_\omega(x)] + \mathbb{E}_{x\sim Q_\theta}[\log(1-D_\omega(x))] \end{equation}\) with \(D_\omega = 1/(1+e^{-V_\omega(x)})\) which corresponds to the following

\(g_f(\nu)= \log(1/(1+e^{-\nu}))\) \(T_\omega(x) = \log (D_\omega(x)) = g_f(V_\omega(x))\) \(f^c(t) = -\log (1-\exp(t))\) \(\log (1-D_\omega(x)) = -f^c(T_\omega(x))\)

Practical considerations in adversarial training

Todo: log trick, DCGAN heuristics

Example divergence and their related functions

WGAN and WAE

Optimal transport (OT)

Kantorovich formulated the optimization target in optimal transport problems as follows

where \(\mathcal P(X\sim P_X, Y\sim P_Y)\) is a set of all join distributions of \((X,Y)\) with marginals \(P_X\) and \(P_Y\).

Wasserstein distance

When \(c(x,y) = \| x-y \| ^p\) for \(p \ge 1\), \(W_c^{1/p}\) is called p-Wasserstein distance.

The optimization problem is highly intractable in general, due to the constraint. However when \(p=1\), Kantorovich-Rubinstein duality holds:

The family of divergences from \(f\)-divergence only consider the relative probability (the ratio between two probability density functions) and do not measure the closeness of the underlying outcomes. With disjoint support or overlapping support but intersections that yield zero measure, the divergence between a target distribution and a \(\theta\)-parameterized distribution might not be continuous with respect to \(\theta\). Wasserstein distance on the other hand does take into account the underlying topology of the outcomes and is continuous and differentiable almost everywhere with respect to \(\theta\) and thus almost always provide useful gradient for optimization.

Wasserstein GAN (WGAN)

Following the dual form of \(W_1\), we can form a generative-adversarial model for a data distribution \(P_D\) and model \(Q_\theta\) with auxiliary function \(f\) that is 1-Lipschitz continuous.

$$ Practical considerations for WGAN

Todo: Gradient clipping with K-Lipschitz constraint on \(f\); Soft gradient penalty (WGAN-GP)

Wasserstein Auto-encoder (WAE)

Rather than working with the dual or Wasserstein distance, which only holds for \(W_1\), we can also work with the primal form directly. As shown in Tolstikhin, Ilya, et al. “Wasserstein Auto-Encoders.” the following holds when we have a deterministic decoder mapping latent variable \(Z\) to \(Y\) through \(y=G(z)\):

The constraint put on \(Q(Z\vert X)\) is that its marginal needs to equal to \(P(Z)\). To have a feasible optimization problem we relax this constraint with the following constraint-free optimization target with a penalty that assess the closeness between \(Q(Z)\) and \(P(Z)\) via any reasonable divergence. This new objective is named penalized optimal transport (POT).

If the divergence between \(P_Z\) and \(Q_Z\) is intractable to directly calculate, we could use generative-adversarial training to approximate it (see \(f\)-GAN).

Note: if decoder is probabilistic instead of deterministic, we would only have \(W_c(P_X, P_G) \le W_c^\dagger(P_X, P_G)\), so we are minimizing an upper bound of the true OT cost.

Thought: the original paper used JS divergence for \(D_Z\), how about we use Wasserstein distance for \(D_Z\).

Todo: discuss connections to AAE paper

Introduction

This post explains Expectation-Maximization (EM) algorithm from scratch in a fairly concise fashion. The material is based on my own notes, which of course come from a variety of great resources online that are listed in the references section.

EM is one of the most elegant and widely used machine learning algorithms but is sometimes not thoroughly introduced in introductory machine learning courses. What is so elegant about EM is that, as we shall see, it originates from nothing but the most fundamental laws of probability.

Many variants of EM have been developed, and an important class of statistical machine learning methods called variational inference also has a strong connection to EM. The core ideas and derivatives of EM find many applications in both classical statistical machine learning and models that involve deep neural networks, making it worthwhile to have an intuitive and thorough understanding of it, which is what this post attempts to provide.

Notation

• Random variables \(X\), probability distribution \(P(X)\)
• Probability density function (PDF) \(p(\cdot)\), evaluated at value \(x\): \(p(X=x)\) with \(p(x)\) as a shorthand
• PDF with parameter \(\theta\) is noted as \(p_\theta(x)\) or equivalently \(p(x\vert \theta)\)
• Expectation of \(f(x)\) according to distribution \(P\): \(\mathbb{E}_{x\sim P}\left[f(x)\right]\)
• A set is noted as \({x_i}\) or calligraphic letter \(\mathcal X\)

Maximum likelihood

Supposed we had data coming from a distribution \(P_D(X)\), and we want to come up with a model for \(x\) parameterized by \(\theta\): \(p(x;\theta)\) or equivalent noted as \(p_{\theta}(x)\) to best approximate the real data distribution. Further assume all the data samples are independent and identically distributed (iid) with \(P_D(X)\).

To find \(\theta\) under a maximum likelihood scheme we do

Motivation for EM

We might encounter situations where, in addition to observed data \({x_i}\), we have missing or hidden data \({z_i}\). It might literally be data that is missing for some reason. Or, more interestingly, it might be due to our modeling choice. We might prefer to have a model with a set of meaningful but hidden variables \({z_i}\) that help explain the “causes” of \({x_i}\). Good examples of this category would be Gaussian (or other kind of) mixture models, and LDA.

Note to myself: examples when we introduces latent variables just for the sake of making the optimization problem easier?

In either case, we will need to have a model for calculating the joint distribution of \(x\) and \(z\), \(p(x,z;\theta)\), which may arise from assumptions (in the case of missing data) or from models of marginal density functions \(p(z; \theta)\) and \(p(x\vert z; \theta)\). In such cases, the log likelihood can be expressed as

Direct maximization of with respect to \(\theta\) might be challenging, due to the summation over \(z\) inside the log. But the problem would be much easier if we knew the values of \(z\). It is simply the original maximum likelihood problem with all data available.

The collection of \(({x_i}, {z_i})\) is called the complete data. Naturally, \({x_i}\) is the incomplete data and \({z_i}\) is the latent data/variable.

Roughly speaking, EM algorithm is an iterative method that let us to guess \(z_i\) based on \(x_i\) (and current estimate of model parameter \(\hat\theta\)). With the guessed “fill-in” \(z_i\) we now have complete data and we optimize the log likelihood \(\ell(\theta)\) with respect to \(\theta\). We thus iteratively improve our guess of latent variable \(z\) and parameter \(\theta\). We repeat this process until convergence.

In slightly more detail, instead of guessing a single value \(z\) we guess the distribution of \(z\) given \(x\), i.e. \(p(z\vert x;\hat\theta)\). then optimize the expected log likelihood for complete data, i.e. \(\sum_i \mathbb{E}_{z \sim p(z\vert x_i;\hat\theta)}\log p_\theta (x_i, z)\), with respect to \(\theta\) which serves as a proxy (lower bound) for the true objective \(\sum_i \log p_{\theta}(x_i)\). Repeat until converge.

(Note in fact guessing a single value for \(z\) is also a valid strategy. It corresponds to a variant of EM and is what we do in the well-known K-means algorithm, where we guess a “hard” label on each data points.)

The nice thing about EM is that it comes with theoretical guarantee of monotonic improvement on the true objective even through we directly work with a proxy (lower bound) of it. Note however the rate of convergence will depend on the problem and the convergence is not guaranteed to be towards the global optima.

Formulation

As before, we start with the log likelihood

Here I switched the summation over \(z\) to integral assuming \(z\) is continuous, just to hint this is a possibility. The last step used Jensen’s inequality and the fact log function is strictly concave. So far we do not have any restrictions on the distribution \(Q\), apart from \(q(z)\) being a probability density function and it is positive where \(p_\theta(x_i,z)\) is.

Using the result above, let’s define the last quantity as \(\mathcal L(q,\theta)\). It is usually called ELBO (Evidence Lower BOund) as it is a lower bound of \(\ell(\theta)\).

Just to reiterate what we have done so far: our goal is to maximize \(\ell(\theta)\), we exchanged the place of the log and integral over \(z\) and got a lower bound \(\mathcal L\).

We can show that the difference between \(\ell(\theta)\) and \(\mathcal L(q,\theta)\) is

where we used the fact Kullback-Leibler (KL) divergence \(D_{KL}\) is defined as

In general, KL divergence is always nonnegative and is zero if and only if \(q(x) = p(x)\). So in our case, the equality \(\ell(\theta) = \mathcal L(q,\theta)\) holds if and only if \(q(z) = p_\theta(z\vert x_i)\). When this happens, we say the bound is tight. In this case, it makes sense to note \(q(z)\) as \(q(z\vert x_i)\) to make the dependence on \(x_i\) clear.

EM algorithm and monotonicity guarantee

The EM algorithm is remarkably simple and it goes as follows.

• E-step (of \(t\)-th iteration):
  • Let \(q^t(z) = p(z \vert x_i; \hat\theta^{t-1})\), which is calculated as shown in Eq. \(\ref{eq:E}\)
  • Due to our particular choice of \(q^t\), at current estimate of \(\hat\theta^{t-1}\) the bond is tight: \(\mathcal L(q^t,\hat\theta^{t-1}) = \ell(\hat\theta^{t-1})\)
• M-step
  • Maximize \(\mathcal L(q^t,\theta)\) with respect to \(\theta\), see Eq. \(\ref{eq:M}\)
  • This step improves ELBO by finding a better \(\theta\): \(\mathcal L(q^t,\theta^t) \ge \mathcal L(q^t,\theta^{t-1})\)

The calculation in E-step is

Just to spell out the function \(\mathcal L(q^t,\theta)\) that we maximize in M-step.

With the preparation earlier we can also easily show the theoretical guarantee on monotonic improvement over the optimization objective \(\ell(\theta)\).

Why the “E” in E-step

By the way, the reason it is called E-step is because in that step we do the necessary calculation to figure out the form of \(\mathcal L(q,\theta)\) as a function of \(\theta\) which we then optimize in the M-step. The form of \(\mathcal L(q,\theta)\) is the expectation of the log likelihood of complete data over the estimated distribution of the latent variable \(z\).

EM as maximization-maximization

Because the particular choice \(q^t(z)\) in E-step is to have diminishing \(D_{KL}(q(z) \| p_\theta(z\vert x_i))\), thus E-step can be viewed as maximizing \(\mathcal L(q,\hat\theta^{t-1})\) with respect to \(q\) and M-step as maximization with respect to \(\theta\). So we are doing alternating maximization on the EBLO with respect to \(q\) and \(\theta\).

This maximization-maximization view offers justification for partial E-step (when the required calculation in exact E-step is intractable) and partial M-step (i.e. only find a \(\theta\) that increases the ELBO rather than maximizes it). Under this view, the direct maximization on ELBO as a goal offers a strong connection to variational inference as will be discussed briefly below.

Example: Gaussian Mixture

In the context of Gaussian Mixture Model (GMM), \(z_i\) associated with \(x_i\) takes the value \({1,2,\dots\,n_{g}}\), where \({n_g}\) is the number of Gaussians in the model. Thus \(z_i\) indicates which Gaussian cluster observed data point \(x_i\) belongs to. The set of parameter \(\theta\) includes those parameterize the marginal distribution of \(z\), \(P(Z;\vec \pi)\). \(\vec \pi = [\pi_1, \pi_2, \dots, \pi_{n_g}]\), with \(\sum_i^{n_g} \pi_i = 1\) and \(\pi_i > 0\). Also, \(\theta\) include those parametrized the conditional distribution of \(P(X \vert Z=z_i; \mu_i, \sigma_i) \sim \mathcal N(\mu_i, \sigma_i)\).

For a detailed walk-through see Andrew Ng’s CS229 lecture notes and video

Variants and extensions of EM

GEM and CEM

A popular variant to EM is that in Eq. \(\ref{eq:M}\) we merely find a \(\hat\theta^t\) that increases (rather than maximizes) \(\mathcal L(q^t,\theta)\). It is easy to see \(\ref{eq:monotone}\) and the monotonicity guarantee still holds in this situation. This algorithm is proposed in the original EM paper and called Generalized EM (GEM).

Another variant is the point-estimate version we mentioned earlier, where instead of having \(q^t(z) = p(z\vert x_i; \hat\theta^{t-1})\) in the E-step, we take \(z\) to be a single value - the most probable one, i.e. \(\hat{z}^t=argmax_z p(z\vert x_i; \hat\theta^{t-1})\) or equivalently taking \(q^t(z) = \delta(z-\hat{z}^t)\). In this case, the integral in \(\ref{eq:M}\) is greatly simplified, but the first equality in \(\ref{eq:monotone}\) does not hold any more and we lose the theoretical guarantee. This algorithm is also called Classification EM (CEM).

Stochastic EM

As we can see in Eq. \(\ref{eq:M}\), we need to go through all data points in order to update \(\theta\), which could be long process for large data sets. In much of the same spirit as stochastic gradient descent, we could sample subsets of data and run the E- and M-step on these mini batches. The same idea can be used for variational inference mentioned below, on the update of global latent variables (such as \(\theta\)).

Variational inference

The computation of the optimal \(q(z)\), i.e. \(q(z) = p(z \vert x_i; \hat\theta_{t-1})\) in E-step might be intractable. Especially, the integral in the denominator of Eq. \(\ref{eq:E}\) does not have closed form solution for many interesting models. In this case we can take the view of EM as maximization-maximization and try to come up with better and better \(q(z)\) to improve the ELBO. In order to proceed with such variational optimization tasks, we need to specify the functional family \(\mathcal Q\) from which we will choose \(q(z)\). Depending on the assumptions a number of interesting algorithms have been developed. The most popular one is probably mean-field approximation.

Note that in a typical variational inference framework, the parameter \(\theta\) is treated as first class variables that we would do inference on (i.e. getting \(p(\theta\vert x)\)) rather than taking a maximum likelihood single point estimation, so \(\theta\) become part of the latent variables and absorbed into the notation \(z\). Thus, \(z\) includes global variables such as \(\theta\) and local variables such as the latent labels \(z_i\) associated with each data point \(x_i\).

In mean-field method the constraint we put on \(q(z)\) is that it factorizes, i.e. \(q(z) = \prod_k q_k(z_k)\). This is saying that all latent variables are mutual independent, by assumption. This seemingly simple assumption brings in remarkable simplifications in the calculation of integrals and especially the expectations of log likelihood involved. It leads to a coordinate ascent variational inference (CAVI) algorithm that allows closed-form iterative calculation for certain model family. The coordinate updates on local variables corresponds to the E-step in EM, while the updates on global variables corresponds to the M-step in EM.

For more about this topic see: D. M. Blei, A. Kucukelbir, and J. D. McAuliffe, “Variational Inference: A Review for Statisticians,” J. Am. Stat. Assoc., vol. 112, no. 518, pp. 859–877, 2017.

References

Todo: add citation in text; for now just core dumped some references here

In no particular order:

1. A. P. Dempster, N. M. Laird, and D. B. Rubin, “Maximum likelihood from incomplete data via the EM algorithm,” J. R. Stat. Soc. Ser. B Methodol., vol. 39, no. 1, pp. 1–38, 1977.

2. R. M. Neal and G. E. Hinton, “A View of the Em Algorithm that Justifies Incremental, Sparse, and other Variants,” Learn. Graph. Model., pp. 355–368, 1998.

3. J. A. Bilmes, “A Gentle Tutorial of the EM Algorithm and its Application to Parameter Estimation for Gaussian Mixture and Hidden Markov Models,” ReCALL, vol. 1198, no. 510, p. 126, 1998.

4. A. Roche, “EM algorithm and variants: an informal tutorial,” pp. 1–17, 2011.

5. M. R. Gupta, “Theory and Use of the EM Algorithm,” Found. Trends® Signal Process., vol. 4, no. 3, pp. 223–296, 2010.

6. M. Jordan, Z. Ghahramani, T. S. Jaakkola, and L. K. Saul, “Introduction to variational methods for graphical models,” Mach. Learn., vol. 37, no. 2, pp. 183–233, 1999.

7. M. J. Wainwright and M. Jordan, “Graphical Models, Exponential Families, and Variational Inference,” Found. Trends® Mach. Learn., vol. 1, no. 1–2, pp. 1–305, 2007.

8. M. Hoffman, D. M. Blei, C. Wang, and J. Paisley, “Stochastic Variational Inference,” vol. 14, pp. 1303–1347, 2012.

9. D. M. Blei, A. Kucukelbir, and J. D. McAuliffe, “Variational Inference: A Review for Statisticians,” J. Am. Stat. Assoc., vol. 112, no. 518, pp. 859–877, 2017.

10. S. Mohamed, “Variational Inference for Machine Learning,” no. February, 2015.

11. Z. Ghahramani, “Variational Methods The Expectation Maximization ( EM ) algorithm,” no. April, 2003.