\documentclass[12pt,letterpaper]{article}
\newcommand\hwnumber{5}
\input{preamble}
\renewcommand{\answer}[1]{}
\renewcommand{\grader}[1]{}

\begin{document}

\textbf{Instructions}: Complete the following questions. There are a total of 29 points for this homework. Each question is marked with its corresponding points. We will be using 
\href{https://colab.research.google.com/drive/1GGyS7tOcTYRGEl3p5NxVHBtpHd3EcLIj?usp=sharing}{this notebook}. Please include any code you write in your report or simply convert your notebook to a PDF and submit that alongside your report.

\textbf{Collaboration: }If you collaborate with other students on the homework, list the names of all your collaborators.
  
\textbf{Submission}: Upload a \textbf{PDF} of your response through \textbf{Canvas} by \textbf{11/12 at 11:59pm}.

\textbf{Notation}: We will use \href{https://phillipi.github.io/6.s898/materials/notation.pdf}{this set of math notation specified on course website}, whose \LaTeX\ source is available on Canvas. For example, $c$ is a scalar, $\vb$ is a vector and $\mW$ is a matrix. You are encouraged (although not enforced) to follow this notation.

\noindent\makebox[\linewidth]{\rule{\linewidth}{0.4pt}}


\subsection*{Variational Autoencoders (14pt)}



Variational autoencoders (VAEs), unlike standard autoencoders, are generative models. We would like to sample a vector from some standard normal distribution $p(\textbf{z})\sim\mathcal{N}(\textbf{0}, \textbf{I})$ and have that sample map to some element in the data distribution $p(\textbf{x})$. 
\begin{figure}[h]
    \centering
    \includegraphics[width=1\linewidth]{figure/vae.pdf}
    \caption{A variational autoencoder (VAE) diagram}
    \label{fig:diffusion}
\end{figure}

Instead of directly maximizing $p(\textbf{x})$, we instead aim to maximize a lower bound on $p(\textbf{x})$ called the evidence lower bound (ELBO) \cite{vaes}:
\begin{equation}
\mathcal{L}_{\text{ELBO}} =  \underbrace{\mathbb{E}_{q_{\phi}(\textbf{z}|\textbf{x})}\big[\text{log}\,p_\theta(\textbf{x}|\textbf{z})\big]}_{\text{first term}}
-\underbrace{D_{KL}(q_{\phi}\big(\textbf{z}|\textbf{x}) \parallel p(\textbf{z})\big)}_{\text{second term}},
\end{equation}
where $p(\textbf{z})=\mathcal{N}(\textbf{0}, \textbf{I})$ is a fixed prior distribution over the latent $\textbf{z}$.

\begin{enumerate}
\item \textbf{(1 pt)} Explain what each of each of the two terms in the above loss function (i.e., the first term and the second term) is doing. Please answer in a few sentences. Note that here we aim to \textit{maximize} the loss function. 


\item \textbf{(1 pt)} In a VAE, each input datapoint $x$ is mapped to a distribution in the latent space by encoding each $x$ as both a mean $\mu$ and variance $\sigma^2$:
\begin{align*}
    \mu_z &= f_\phi^\mu(x), \\
    \sigma^2_z &= f_\phi^\sigma(x). 
\end{align*}
In this question, we will think about how rich the structure of the means in the embedding space (i.e., the quantities $f_\phi^\mu(x)$ for all input datapoints) alone must be in a fully-trained VAE.
Suppose you successfully train a VAE such that $q_{\phi}(\textbf{z})$ becomes a unit Gaussian, and $p_\theta(x)$ is equal to the data distribution $p_\text{data}$. In other words, assume you achieved \textit{perfect} encoder and decoder after training a VAE. 







\begin{enumerate}

    \item \textbf{(0.5 pt)}  Does this imply that decoding only the means of the embeddings recovers the data distribution? In other words, are the sets of means of embedding sufficient for data distribution recovery, or are the variances of embeddings also necessary?
    Explain your answer in a few sentences.
    

    \item \textbf{(0.5 pt)}  Does this imply that the means of the embeddings (i.e., $f_\phi^\mu(x)$ above) are Gaussian distributed? In other words, is the random variable $f_\phi^\mu(x)$ Gaussian when $x \sim p_{\text{data}}$? Explain your answer in a few sentences.

    \hint{Note that the random variable $\textbf{z} = \mathcal{N} (f_\phi^\mu(x), f_\phi^\sigma(x))$ is unit Gaussian by assumption, when $x \sim p_{\text{data}}$. Does this imply that its mean is also Gaussian?}



\end{enumerate}

\item \textbf{(1 pt)} Suppose you successfully train an autoencoder (not VAE) such that $g(f(\textbf{x}))=\textbf{x}$ where $f$ is the deterministic encoder and $g$ is the deterministic decoder. Does this imply that $g(\textbf{z}) \sim p_{\texttt{data}}$ when $\textbf{z} \sim \mathcal{N}(\textbf{0}, \textbf{I})$? Explain your answer in a few sentences.


\item \textbf{(6 pts)} Implement the VAE architecture in the provided colab. 

\item \textbf{(1 pt)} Train the VAE model on the provided FashionMNIST dataset.

\item \textbf{(4 pt)} Using your trained model, complete the $\texttt{plot\_latents}$ function such that for a given pair of latent dimensions $(i, j)$ the $\texttt{plot\_latents}$ plots a $10\times{10}$ grid of images sampled from different pairs of latent values in dimensions $i, j$ (Figure 2). Latent dimensions not equal to $i$ or $j$ should be set to zero.

\begin{figure}[h]
    \centering
    \includegraphics[width=0.3\linewidth]{figure/latents.png}
    \caption{Example latent visualization.}
    \label{fig:seq}
\end{figure}


\end{enumerate}


\vspace{-0.2in}
\subsection*{Diffusion Models (15pt)}
A diffusion process works by sequentially adding noise to a given input $x_0$ for $T$ timesteps, and our goal is to learn a function that removes the added noise (i.e., \textit{denoising}). While there are many variations around this central concept, we will be looking closely at one particular approach from \citep{diffusion}.

As $T$ approaches infinity, this process results in noise drawn from a zero-mean isotropic Gaussian at $\textbf{x}_T$. This process is depicted in Figure~\ref{fig:diffusion}. The random variable $\textbf{x}_t$ conditioned on $\textbf{x}_{t-1}$ is distributed according to $q(\textbf{x}_t|\textbf{x}_{t-1})$ with the following form:
\begin{equation}
q(\textbf{x}_t|\textbf{x}_{t-1}) = \mathcal{N}(\sqrt{\alpha_t}\,\textbf{x}_{t-1},{\beta}_t\textbf{I}),  \;\;\text{where}\;\;  \alpha_t = 1-{\beta_t}.
\end{equation}

\textit{Note: the distribution $q$ refers to the joint distribution over over all $\mathbf{x}_0, \ldots, \mathbf{x}_T$. Conditionals of this distribution are denoted as $q(\cdot|\cdot)$ for subsets of the variables in the joint configuration.}

Let's first try to understand the properties of the forward diffusion process in terms of $q(\textbf{x}_{t}|\textbf{x}_0)$. We will use these properties later to derive a loss for the reverse diffusion process.

\begin{enumerate}[resume]
    \item \textbf{(2pt)} Adding two independent Gaussian random variables results in another Gaussian random variable. Using this property, one can derive a closed-form representation 
    %of $\mathcal{N}(\cdot, \cdot)$ 
    for the Gaussian distribution  $q(\textbf{x}_{t}|\textbf{x}_0)$ in terms of $\alpha_t$ and $\textbf{x}_0$. 
    Prove that 
    \begin{align}
        q(\textbf{x}_{t}|\textbf{x}_0) = \mathcal{N}(\sqrt{\bar{\alpha}_t}\textbf{x}_0, (1-\bar{\alpha}_t)\textbf{I}), \label{eqn:q_xt_x0}
    \end{align} 
    where $\bar{\alpha}_t$ is referred to  as $\prod_{s=1}^t{\alpha_s}$.
    Provide a detailed explanation of how you arrived at your answer. 
    
\end{enumerate}

\begin{enumerate}[resume]
    \item \textbf{(1pt)} Let us calculate how far $\textbf{x}_T$ is from being $\mathcal{N}(\textbf{0}, \textbf{I})$ in KL divergence. For any $d$-dimensional Gaussian distribution $\mathcal{N}(\textbf{a}, \Sigma)$, we have
    \begin{align*}
        D_{\text{KL}}\big( \mathcal{N}(\textbf{a}, \Sigma) \parallel \mathcal{N}(\textbf{0}, \textbf{I})\big) = \frac{1}{2} \left( \| \textbf{a}\|_2^2 + \operatorname{tr}(\Sigma) - \log(\det(\Sigma))  - d \right),
    \end{align*}
    where $\operatorname{tr}(\Sigma)$ is the sum of diagonal entries of the covariance matrix $\Sigma$. Use the above formula and compute  $D_{\text{KL}}\big( q(\textbf{x}_{T}|\textbf{x}_0) \parallel \mathcal{N}(\textbf{0}, \textbf{I})\big)$. Conclude that if $T$ is large, then $q(\textbf{x}_{T}|\textbf{x}_0)$ is close to being unit Gaussian. 


\end{enumerate}

Now, our real goal is to learn a function $p_\theta(\textbf{x}_{t-1}|\textbf{x}_t) = \mathcal{N}(\mu_\theta(\textbf{x}_t), \Sigma_t)$ that removes the noise added (i.e.,  \textit{denoising}). Indeed, learning such a function enables us to remove noise from $\mathbf{x}_t$ to recover $\mathbf{x}_{t-1}$. By starting with a completely noisy image (i.e., a unit Gaussian), we can perform sequential denoising to generate an image with this approach.

One strategy to learning $p_\theta(\textbf{x}_{t-1}|\textbf{x}_t)$ is to minimize the KL-divergence between this distribution and the true reverse process $q(\textbf{x}_{t-1}|\textbf{x}_t)$:

\begin{figure}[h]
    \centering
    \includegraphics[width=0.8\linewidth]{figure/diffusion_figure.png}
    \caption{The directed graphical model for the diffusion process.}
    \label{fig:diffusion}
\end{figure}

But what is the form of the reverse diffusion process $q(\textbf{x}_{t-1}|\textbf{x}_t)$? We will derive this in the following questions.





    




To derive this expression, we use the closed form for the expression $q(\textbf{x}_{t-1}|\textbf{x}_t, \textbf{x}_0)$, which works out to be (you may refer to \citet{diffusion} for more details):
\begin{equation}
q(\textbf{x}_{t-1}|\textbf{x}_t, \textbf{x}_0) = \mathcal{N}(\tilde{\mu}(\textbf{x}_t, \textbf{x}_0), \tilde{\beta}_t \textbf{I})
\end{equation}
\begin{equation}
\tilde{\beta}_t = \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}}\cdot\beta_t
\end{equation}
\begin{equation}
\tilde{\mu}(\textbf{x}_t, \textbf{x}_0) = \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_{t}}\cdot{\textbf{x}_t}+\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_{t}}\cdot{\textbf{x}_0}
\end{equation}
Also, note from Eqn. \ref{eqn:q_xt_x0}, we have that
\begin{align}
        \textbf{x}_t = \sqrt{\bar{\alpha}_t}\textbf{x}_0 +  \sqrt{(1-\bar{\alpha}_t)} \epsilon_t, \label{eqn:xt_in_terms_of_xo}
\end{align}
where $\epsilon_t\sim\mathcal{N}(\textbf{0}, \textbf{I})$ is a Gaussian random variable with unit variance. 

Using the above expressions:

\begin{enumerate}[resume]

    \item \textbf{(1pt)} Derive an expression for mean and the variance of $q(\mathbf{x}_{t-1} | \mathbf{x}_t)$ in terms of the above quantities. 


    \item \textbf{(1pt)} Derive an expression for $D_{\text{KL}}(q(\textbf{x}_{t-1}|\textbf{x}_t) \parallel p_\theta(\textbf{x}_{t-1}|\textbf{x}_t))$.

    \hint{You can use the following formula for the KL divergence between two $d$-dimensional Gaussian vectors:
    \begin{align*}
        D_{\text{KL}}\left( \mathcal{N}\left(\textbf{a}, \Sigma_a\right) 
        \parallel
            \mathcal{N}\left(\textbf{b}, \Sigma_b\right)    
        \right) = &\frac{1}{2} \Big(
        \log(\det(\Sigma_b)) - \log(\det(\Sigma_a)) - d \\&
                +(\textbf{a} - \textbf{b})^T \Sigma_b^{-1}(\textbf{a} - \textbf{b})   + \operatorname{tr}(\Sigma_b^{-1}\Sigma_a)
          \Big).
    \end{align*}
    }

    
\end{enumerate}

Congrats! You now have an objective function for denoising, which you can use to train a diffusion model.


In practice, however, people often reparameterize the optimization problem in terms of $\epsilon_t = \mathbf{x}_t - \mathbf{x}_{t-1}$ and predict $\epsilon_t$ with a model $\epsilon_{\theta}(\mathbf{x}_t, t)$. Doing this, and other simplifications (see \cite{diffusion} for details), we get the following diffusion modeling loss, which is used in the coding section of this pset:
\begin{equation}
\mathcal{L}_{\texttt{Diffusion}} = \| \epsilon_t-\epsilon_{\theta}(\textbf{x}_t, t)\|^2_2.\label{eq::loss}
\end{equation}
    
    
\begin{enumerate}[resume]
    
    \item \textbf{(6pt)} In the provided colab, complete the $\texttt{FIXME}$s and train the diffusion model.
    
    \item \textbf{(2pt)} Create a plot with six horizontally-aligned subfigures showing the diffusion model inference $\textbf{x}_t$ for $t\in{\{200, 100, 50, 20, 10, 0\}}$. An example can be seen in Figure~\ref{fig:inference}. Your plots will look different.


    
    \begin{figure}[t]
    
    \centering
    \includegraphics[width=0.8\linewidth]{figure/fnmnist.png}
    \caption{Diffusion inference figure example}
    
    \label{fig:inference}
    \end{figure}



    
    \item \textbf{(2pt)} There are many similarities between VAEs and diffusion models. Answer each question with a sentence or two.
    \begin{enumerate}
        \item{ \textbf{(1pt)} What is the equivalent of the encoder, decoder, and latent variable in the diffusion model?} 

        
        \item{ \textbf{(1pt)} What is the difference between the encoder in a diffusion model and VAE?}
        
        
    \end{enumerate}
\end{enumerate}


\bibliographystyle{plainnat}
\bibliography{reference}

\end{document}
