\documentclass[12pt,letterpaper]{article}
\newcommand\hwnumber{4}
\input{preamble}

\begin{document}

\textbf{Instructions}: Complete the following questions. There are a total of 25 points for this homework. Each question is marked with its corresponding points. Some questions are not graded, including all bonus questions. But you are encouraged to think about and attempt them.

If a short ($\leq$ 1 sentence) answer is requested, such as a short expression or comment, please write it concisely in a clearly identifiable location. \fbox{Box} your answers for such problems, especially if handwriting them, (\texttt{e.g. using "$\backslash$fbox$\{$Box$\}$"} in \LaTeX). In addition, \fbox{box} your answers for all ``yes/no'' or multiple choice questions.

\textbf{Please adhere to each questions expected deliverables and answer length!}

For all plotting questions, do not hand-draw the functions, use plotting tools instead.

\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/5 at 11:59PM}. Please complete the page assignment to each section within Gradescope,
otherwise, we will apply a \textbf{3\%} penalty to the homework. This can be done after you upload
your submission PDF.

\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.

\textbf{Note on use of AI Assistants}: This assignment contains a few coding questions. We understand that AI assistants can complete many coding based problems. Therefore, we would like to provide a reminder of our AI assistant policy. You may not simply ask an AI assistant to complete you code for you. This includes Google Colab autocomplete, which should not be used for completing assignments. Our full AI assistant policy can be found \href{https://phillipi.github.io/6.7960/#AI_policy}{here}.  

\textbf{Note on Colab Autocomplete Usage}:
Please disable Colab's code autocomplete for this problem set. Navigate to Tools/Settings/Editor, and toggle off the option labeled "Automatically trigger code completions." For more details, refer to the instructions \href{https://stackoverflow.com/questions/63696360/google-colab-how-to-turn-off-suggestion-window}{here}.


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

This homework has two sections, roughly of equal work. \begin{itemize}
    \item \textbf{(13pt)} The first section focuses on similarity-based learning in both supervised and self-supervised learning, with a focus on the latter. We derive certain properties of such learned representations, and highlight a connection between the supervised objective and the self-supervised objective.
    \item \textbf{(12pt)} The second section focuses on empirically analyzing learned representations from reconstruction-based and similarity-based objectives. Via nearest neighbors, we show how choices of self-supervised objectives can affect the information captured in representations.
\end{itemize}

\vspace{-0.2in}

\vspace{-0.2in}
\subsection*{Similarity-Based Learning: Self-Supervised and Supervised  (13pt)}
Recall from lecture that contrastive learning optimizes an encoder network on certain similarity measures. In this section, we will explore such similarity-based objectives, and see how they relate to contrastive learning and supervised learning.

\textbf{Setting:} For some data space $\mathcal{X}$, we consider data coming from a (equal) mixture of  $N$ distributions $\{p_i\}_{i=1}^N$, where each $p_i$ is a distribution over $\mathcal{X}$, and $N$ is finite. These represent $N$ groups of data samples, where samples within each group are considered \emph{similar}.  
Essentially, $\{p_i\}_{i=1}^N$ defines the similarity relation that we will learn from. 

\textbf{Contrastive Learning:} Consider an encoder $f \colon \mathcal{X} \rightarrow \Sp^{d-1}$, where $\Sp^{d-1}$ is the $(d-1)$-dimensional unit hypersphere $\subset \R^d$ (where the normalization removes a dimension). Usually, $f$ is a deep encoder network computing a vector $\in \R^d$, followed by an $l_2$-normalization step (dividing the vector by its norm). $\Sp^{d-1}$ is our representation space. We consider the following specific form of contrastive loss: \begin{equation}
    \mathcal{L}_\mathsf{contr}(f) \trieq \frac{1}{N} \sum_{i=1}^N 
    \expsubst{x_i \sim p_i\\\forall j \in \{1, \dots, N\},~y_j \sim p_j} \left[ - \log \frac{e^{f(x_i)\T f(y_i)}}{\sum_{j=1}^N e^{f(x_i)\T f(y_j)}} \right],\label{sim:contr-loss}
\end{equation}
where $(x_i, y_i)$ is the \emph{positive pair} and, for $j \neq i$, $(x_i, y_j)$ are the \emph{negative pairs}.

\fbox{%
\parbox{\linewidth-5pt}{%
An example with $N=4$ mixtures:

\includegraphics[width=\linewidth, trim=460 400 485 130,clip]{figure/hw4_contrastive_mixture.pdf}}}

For this section, we consider optimizing over all possible functions. In practice, $f$ is usually a parametrized encoder using deep neural networks. 

\fbox{ %
\parbox{\linewidth-5pt}{%
\textbf{Remarks on our setting (not required for completing the questions):} Our setting and the contrastive loss form in \Cref{sim:contr-loss} may look a bit different from what you encounter elsewhere.  For example, we assume that \begin{itemize}[itemsep=-1.5pt,topsep=3pt]
    \item $N$, the number of similarity groups ($p_i$'s), is finite, 
    \item The logits are not scaled by some temperature parameter
    \item For an element $x_i \sim p_i$. the negatives are always formed from exactly one sample from each other similarity group $p_j$ ($j \neq i$), and never from the same group $p_i$.
\end{itemize} These assumptions make our analysis easier. While the exact results we derive may be slightly different from other analyses, the core messages are the same.
}}
\begin{enumerate}[resume]
    \item \textbf{(1pt; Augmentations and Image Contrastive Learning)} For contrastive learning on images, recall that we often choose positive pairs as augmentations of the same image, and negative pairs as augmentations of different images.  
    
    Consider an image dataset $\{\mathsf{img}_i\}_{i=1}^K$, and $p_\mathsf{aug}({}\cdot{} \mid \mathsf{img})$ the distribution over augmented versions of an image $\mathsf{img}$. 
    
    To cast this into our setting, write down formulas for $N$ and $\{p_i\}_i$ in terms of $\{\mathsf{img}_i\}_{i=1}^K$, $K$, and $p_\mathsf{aug}$. 
    
    \item \label{sim:ins-dis} \textbf{(1pt; Instance Discrimination/Classification)} \Cref{sim:contr-loss} is essentially a cross-entropy loss. For $x_i$ and $\{y_j\}_{j=1}^N$, write down the number of classes $C$, the logits vector $\in \R^C$, and the target class $\in \{1, \dots, C\}$. Are entries of the logits vector bounded within some range? If so, write down the range.
    
    \item \textbf{(0.5pt; Dot Products and Distances)} { Derive the relationship between the Euclidean distance $\lVert u - v \lVert_2$ and the dot product  $u^Tv$ of two vectors $u, v$. Further, show how the cosine similarity between the two vectors relates to the Euclidean distance when both vectors lie on the unit sphere $u, v \in \Sp^{d-1}$.}
    
    \item \label{sim:contr-clf-acc} \textbf{(7pt; From Classification to Representation Geometry)} In Question~\ref{sim:ins-dis}, we cast the contrastive loss as a classification loss, where each combination of $x_i$ and $\{y_j\}_{j=1}^N$ gives a classification question for the encoder $f$. Unlike the usual supervised classification training, the target label here is constructed differently, and the logits follow a weird form involving $f$. 
    
    Let's build some intuition on what this  classification task even means.  In particular, we characterize what encoders $f$ can achieve perfect classification \emph{accuracy}.
    
    \textbf{Assumption:} \begin{equation}
        \text{Each $p_i$ is restricted to a finite set $S_i$, where $p_i(s) >0$ for all $s \in S_i$.} \label{sim:fin-asumpt}
    \end{equation}
    
    \begin{enumerate}
    
        \item \label{sim:fin-ex} \textbf{(2pt)} Suppose $N = 3$, each $p_i$ is a uniform distribution over 4 elements, and $f$ encoder simply outputs a 2-dimensional unit vector (on $\Sp^1$).
    
        In the four plots below, we show four different feature distributions from  different $f$ choices. \textbf{For each one, answer ``Yes'' or ``No'':} does $f$ achieves a perfect classification accuracy (over all choices of $x_i$ and $\{y_j\}_{j=1}^N$)?
        
        \includegraphics[width=\linewidth, trim=0 400 0 300, clip]{figure/hw4_contrastive_feature_distn.pdf}
        
        (Recall that a prediction is correct if the largest entry of $\mathsf{logit}$ is \emph{uniquely} $\mathsf{target}$.)
                
        \item \textbf{(1pt; Feature Geometry)} Recall that we have Assumption~\eqref{sim:fin-asumpt}. For example, the scenario in Question~\ref{sim:fin-ex} is a case satisfying the assumption.
        
        Consider the following quantities: \begin{align}
            \mathrm{diameter}_i(f) & \trieq \max_{s, t \in S_i} \lVert f(s) - f(t) \rVert_2  &\mathrlap{\forall i} \\
            \mathrm{margin}_{i\rightarrow j}(f) & \trieq \min_{\substack{s_i \in S_i \\ s_j \in S_j}}\lVert f(s_i) - f(s_j) \rVert_2  &\mathrlap{\forall i, j}
        \end{align}
        
        These are metrics closely related to the positive and negative pair distributions. 
        In particular, $\mathrm{diameter}_i$ is defined with the positive pairs from $p_i$ (supported on set $S_i$), and $\mathrm{margin}_{i\rightarrow j}$ is defined with the negative pairs from $p_i$ and $p_j$.
        
        Describe in approx. 1-2 lines what $\mathrm{diameter}_i(f)$ and $\mathrm{margin}_{i\rightarrow j}(f)$ geometrically represent in terms of the distance between feature vectors.
        
        \item  \label{sim:dia-margin} \textbf{(1pt; Optimal Feature Geometry)} Multiple choice: Which of the following conditions on encoder $f$ ensures a perfect accuracy? Choose one:
         \begin{enumerate}
             \item \begin{equation*}
                 \forall i,\quad  \exists j \neq i\quad \textsf{ such that }\quad \mathrm{diameter}_i(f) < \mathrm{margin}_{i\rightarrow j}(f)
             \end{equation*}
             
             (\ie, for all distributions $i$, its diameter is smaller than the margin between $i$ and \textbf{some} other $j$.)
             
             \item \begin{equation*}
                 \forall i \neq j,\quad \mathrm{diameter}_i(f) < \mathrm{margin}_{i\rightarrow j}(f)
             \end{equation*}
             
             (\ie, for all distributions $i$, its diameter is smaller than the margin between $i$ and \textbf{all} other $j$.)
             
             \item  \begin{equation*}
                 \exists i^*\quad \textsf{ such that }\quad\mathrm{diameter}_{i^*}(f) < \min_{i \neq j} \mathrm{margin}_{i\rightarrow j}(f)
             \end{equation*}
             
             (\ie, the diameter of $i_*$ is smaller of any margin between two different distributions.)
             
         \end{enumerate}
        
        \item {(Optional) Question: Explain why, as the dimensionality $d$ of the feature space increases, the margin between different classes can become smaller. How does this affect the choice of encoder function $f$?}\\
        
        \item \label{sim:align-margin-unif} \textbf{(2pt; Alignment, Invariance, Margin, and Uniformity)} Consider the case where $p_i$'s have disjoint supports (\ie, $S_i$'s). 
        Answer the following questions: \begin{enumerate}
            \item \textbf{(0.5pt; Alignment)} To more easily satisfy the condition you obtained in Question~\ref{sim:dia-margin}, should we have larger or smaller $\mathrm{diameter}_i$'s? Answer ``larger'' or ``smaller''. 
                        
            \item \textbf{(0.5pt; Margin)} To more easily satisfy the condition you obtained in Question~\ref{sim:dia-margin}, should we have larger or smaller $\mathrm{margin}_{i\rightarrow j}$'s? Answer ``larger'' or ``smaller''. 
                        
            \item \textbf{(1pt)} Consider an optimal encoder $f^*$ that extremizes $\mathrm{diameter}_i$'s and $\mathrm{margin}_{i\rightarrow j}$'s based on your answers above.  \begin{enumerate}
                \item \textbf{(0.5pt; Invariances)}
                We say that an encoder $f$ is \emph{invariant} to variations within a set $S \subset \mathcal{X}$ if $f(x)$ is constant $\forall x \in S$.
                
                What are the \emph{invariances} that $f^*$ learn (if any)? Provide your answer in the form of a collection of \textbf{sets} $\subset \mathcal{X}$, where $f$ is invariant to variations within each, and these sets are ``largest'' in the sense that $f^*$ is \textbf{not} invariant to the union of any two of them.
                
                \item \textbf{(0.5pt; Uniformity)}
                In geometry, a well known problem is to find max-margin placements of $N$ unit-vectors on $\Sp^{d-1}$ \citep{tammes1930origin}, whose solution is ``roughly uniformly'' placing the points $\Sp^{d-1}$.
                
                Consider the above Tammes' problem and the problem of extremizing $\mathrm{diameter}_i$'s and $\mathrm{margin}_{i\rightarrow j}$. Discuss the relation between them in 1-3 sentences.
                
            \end{enumerate} 
        \end{enumerate}
        
        \item \textbf{(1pt; Overlapping Distributions)} \textbf{Without the assumption that} $p_i$'s have disjoint supports (\ie, $S_i$'s). Answer the following questions: \begin{itemize}
            \item \textbf{(0.5pt)} Answer ``Yes'' or ``No'': Can the condition you obtained in Question~\ref{sim:dia-margin} be satisfied with possible overlapping supports?
                        
            \item \textbf{(0.5pt)} What will happen to two $p_i$ and $p_j$ that have overlapping support: \begin{equation}
                \mathrm{supp}(p_i) \bigcap \mathrm{supp}(p_j) = S_i \bigcap S_j \neq \varnothing, \mathrlap{\qquad i \neq j}
            \end{equation}
            
            Provide 1-2 sentence description on how their features are placed by an optimal encoder $f$. 
            
            Your answer doesn't need to provide an exact characterization, but should mention the distance between feature vectors from these two distributions.

        \end{itemize}
    \end{enumerate}
    
    \textbf{Remark:} This accuracy-based analysis already reveals how contrastive learning groups features together based on the positive pair distributions. In particular, Question~\ref{sim:align-margin-unif} says that contrastive learning might prefer tight clusters that are uniformly placed in the representation space. While many assumptions are used to obtain these insights, a more involved analysis shows that they hold true generally \citep{wang2020hypersphere}. 
    
    \item \label{sim:perf-acc-scale} \textbf{(2pt; Normalization in Contrastive Learning)} Consider an encoder $f$ with perfect accuracy on the data distributions (as defined in Question~\ref{sim:contr-clf-acc}). Answer the following questions: \begin{itemize}
        \item \textbf{(1pt)} Answer ``Yes'' or ``No'': Does $f$ achieve zero contrastive loss (\Cref{sim:contr-loss})?
        \item \textbf{(1pt)} Suppose we relax the constraint that encoders must output to $\Sp^{d-1}$, and instead allow them to output any value $\in \R^d$. Consider optimizing such encoders to minimize the contrastive loss.
        \begin{itemize}
            \item \textbf{(0.5pt)} Is there an encoder that attains arbitrarily small contrastive loss?
            
            \hint{One approach, temperature scaling divides the logits by a scalar parameter $softmax = e^{z/T} / \sum_i e^{z_i/T}$. Think about the effect of logit scaling on cross-entropy classification loss. For more information on temperature scaling, see \url{https://arxiv.org/pdf/1706.04599.pdf} }
                        
            \item \textbf{(0.5pt)} Are there any potential numerical stability issues? 
            
        \end{itemize}
    \end{itemize}
    
    \item \label{sim:clf-via-contr} \textbf{(1.5pt; Cross-Entropy Supervised Learning as Dual-Encoder Contrastive Learning)}

    In fact, the familiar cross-entropy supervised classifier learning can be exactly cast as a Dual-Encoder contrastive learning method. 
    
\begin{mdframed}[innertopmargin=0.1cm,innerleftmargin=0.1cm,innerrightmargin=0.1cm]
    \textbf{Dual-Encoder Contrastive Learning}: Consider \textbf{two data domains} $\mathcal{X}_1$ and $\mathcal{X}_2$ (\eg, images and text), \textbf{two encoders} \begin{align}
        f_1 &\colon \mathcal{X}_1 \rightarrow \R^{d} \\
        f_2 &\colon \mathcal{X}_2 \rightarrow \R^{d},
    \end{align} and positive pair distributions $\{p_i\}_{i=1}^N$, where each $p_i$ is a distribution over the \textbf{product space} $\mathcal{X}_1 \times \mathcal{X}_2$.  Define the following form of \textbf{Dual-Encoder} contrastive loss: \begin{equation}
    \mathcal{L}_\mathsf{dual\mhyphen enc\mhyphen contr}(f_1, f_2) \trieq \frac{1}{N} \sum_{i=1}^N 
    \expsubst{\forall j \in \{1, \dots, N\},\\(x^{(1)}_j, x^{(2)}_j) \sim p_j} \left[ - \log \frac{e^{f_1(x_i^{(1)})\T f_2(x_i^{(2)})}}{\sum_{j=1}^N e^{f_1(x_i^{(1)})\T f_2(x_j^{(2)})}} \right],\label{sim:x-contr-loss}
\end{equation}
where $(x^{(1)}_i, x^{(2)}_i)$ is the \emph{positive pair} and, for $j \neq i$, $(x^{(1)}_i, x^{(2)}_j)$ are the \emph{negative pairs}.\\[-0.75em]

This is referred to as \emph{cooccurrence-based} contrastive learning in notes.\\[-0.75em]

\textbf{Remark :} We use unconstrained encoder (output $\in \R^d$) to make connections precise. In practice, dual-encoder methods usually still use normalized outputs $\in \Sp^{d-1}$. 

\end{mdframed}

    Assume a labelled dataset for classification $\mathcal{D}_\mathsf{labelled} \trieq \{(x_k, y_k)\}_{k=1}^{K}$, where $x_k \in \mathcal{X}$ are data samples, and $y_k \in \{1, 2, \dots, C\}$ are labels from $C$ evenly-represented classes.
    
    With standard supervised training, consider optimizing a \textbf{deep classifier} $g \colon \mathcal{X} \rightarrow \R^C$ \wrt cross-entropy loss on $\{x_k, y_k\}_{k=1}^K$: \begin{equation}
        \mathcal{L}_\mathsf{xce}(g) \trieq \expsub{x_k, y_k \sim \mathcal{D}_\mathsf{labelled}} \bigg[ -\log \underbrace{\frac{e^{g(x_k)[y_k]}}{\sum_c e^{g(x_k)[c]}}}_{\textsf{softmax}(g(x_k))[y_k]}\bigg].\label{sim:sup-xce}
    \end{equation}
    
    Let's refer to this training procedure as \textbf{supervised cross-entropy classification}, where we say that $g$ correctly predicts on pair $(x, y)$ iff \begin{equation}
        \argmax_{c \in \{1, 2, \dots, C\}} \underbrace{g(x)}_{\mathclap{\mathsf{logits}}}[c] = y,\mathrlap{\text{\quad where $\argmax$ is unique.}}
    \end{equation}
    
    Compare it with the \textbf{dual-encoder contrastive learning} procedure, and answer the following questions in write-up: \begin{itemize}
        \item \textbf{(1pt)} Provide the exact forms of $N$, $\{p_i\}_i$, $d$, $\mathcal{X}_1$, $\mathcal{X}_2$, $f_1$, and $f_2$, so that training $g$ \wrt $\mathcal{L}_\mathsf{xce}$ exactly corresponds to training $f_1$ and $f_2$ \wrt $\mathcal{L}_\mathsf{dual\mhyphen enc\mhyphen contr}(f_1, f_2)$. 
            
        \hint{
        Think about how to structure the cross entropy loss as a dual encoder contrastive loss problem.
        As an example of an a possible alternate answer, you may decompose $g(x) = W\T [g_\mathsf{feature}(x) :: 1]$, where $::$ denotes concatenation, $W$ is the parameters of the final linear layer in $g$, and $g_\mathsf{feature}$ is the computation of $g$ before the final layer with output $\in \R^{d_\mathsf{feature}}$.
        \begin{align*}
            N & = C \\
            p_i & = \mathrm{Uniform}\{x_k \colon y_k = i\}\\
            d &= d_\mathsf{feature} + 1 \\
            \mathcal{X}_1 &=\mathcal{X} \\
            \mathcal{X}_2 &=\{1, 2, \dots, C\} \\
            f_1(x) & = [g_\mathsf{feature}(x) :: 1] \in \R^{d_\mathsf{feature} + 1} \\
            f_2(c) & = W_c \in \R^{d_\mathsf{feature} + 1}.
        \end{align*}
        }
        
        \item \textbf{(0.5pt)} Show that under this view, the supervised cross-entropy classification accuracy of $g$ coincides with the contrastive loss accuracy defined in Question~\ref{sim:contr-clf-acc}.
        
    \end{itemize}
    
\end{enumerate}



\newpage
\vspace{-0.2in}
\subsection*{Reconstruction and Similarities in Representation Learning (12pt)}
We consider performing unsupervised representation learning on a dataset of colored shapes.  Each image has resolution  64×64 . Across the dataset, there are three underlying varying factors, as shown in \Cref{fig:shapes-variations}.


\begin{figure}[h]
    \centering
    \includegraphics[width=\linewidth, trim=150 50 120 45, clip]{hw4_dataset_viz/uncond.pdf}\vspace{-10pt}
    \caption{Random unconditional samples of the dataset.}
    \label{fig:shapes-uncond}
\end{figure}



\begin{figure}[h!]
    \centering
    \includegraphics[width=\linewidth, trim=280 410 630 280, clip]{hw4_dataset_viz/variations.pdf}\vspace{-10pt}
    \caption{Three types of variations in the dataset.}
    \label{fig:shapes-variations}
\end{figure}


We will use a finite training dataset of size $N$, denoted as \begin{equation}
    \mathcal{D}_\mathsf{train} \trieq \{x_i\}_{i=1}^N,
\end{equation}where each $x_i \in \mathcal{X}$ is such an image. The goal of representation learning is to optimize an encoder function $f \colon \mathcal{X} \rightarrow \R^d$, so that $f(x)$ is a good representation/latent/embedding/encoding of $x$. Here $\R^d$ is called \emph{the representation space}, with usually a small $d$ aiming for a compact representation of the dataset that captures important features.


\textbf{Instructions:} We will use \href{https://drive.google.com/file/d/1DsXwdehCKs9j3NtU5_Ly3cZM1tRSlSfv/view?usp=sharing}{this notebook}. There are only a few \hltexttt{FIXME} place where you need to fill in and submit the code. Majority of the code are already provided. You are not required to fully read through the implementation of them, but should conceptually understand what the functions/classes with \textbf{docstrings} are doing. You will be asked to attach certain plots, and provide certain discussions about the results in your final write-up.  The training time for some contrastive encoders can be close to 1 hour, while the autoencoder and some other contrastive encoders are much faster. \textbf{Please include all code as a PDF attached to the end of your HW submission.}

\begin{enumerate}[resume]
    \item \textbf{(6pt; Autoencoder)}
    Autoencoder objective involves two functions: \begin{align}
        f &\colon \mathcal{X} \rightarrow \R^d \tag{encoder} \\
        g &\colon \R^d \rightarrow \mathcal{X}  \tag{decoder}.
    \end{align}
    The objective is reconstruction, saying that $g$ should be able to fully reconstruct input data $x$, by only looking at the representation $f(x)$: \begin{equation}
        \mathcal{L}_\mathsf{AE}(f; g) = \frac{1}{N} \sum_{i=1}^N \underbrace{\mathsf{MSE}(g(f(x)), x)}_{\vphantom{2^{2^{2}}}\mathclap{\textsf{mean squared error between reconstruction and original $x$}}}.\label{rep:ae}
    \end{equation}
    \begin{enumerate}
        \item \textbf{(3pt; Theoretical Optimum)}
        Consider two functions $f^*$ and $g^*$. All we know is that they perfectly solve the Autoencoder objective in \Cref{rep:ae} on the finite dataset $ \mathcal{D}_\mathsf{train} \trieq \{x_i\}_{i=1}^N$, where all samples are assumed to be \textbf{distinct} and samples from some $p_x$. 
        
        Answer each of the following questions with respect to the shape dataset in \Cref{fig:shapes-variations} with one of \{``Yes'', ``No'', ``Maybe''\} and briefly explain your answer (approx. 1 sentence).  \begin{enumerate}
            \item \textbf{(0.5pt)} For some $x_i \in  \mathcal{D}_\mathsf{train}$, can its dataset index $i$ be decided by $f^*(x_i)$ alone? 
            
            (A property of $x$ can be decided by $f^*(x)$ alone if there exists a function $h$ such that the property of $x$ can be written as $h(f^*(x))$.)

            \item \textbf{(0.5pt)} For some $x_i \in  \mathcal{D}_\mathsf{train}$, can the color of $x_i$ be decided by $f^*(x_i)$ alone? 
            
            \item \textbf{(0.5pt)} For three $\{x, y, z\} \subset  \mathcal{D}_\mathsf{train}$, where $x$ and $y$ are both red, and $z$ is blue. Compared to  $x$ and $z$, are $x$ and $y$ closer in representation space?
            
            \Ie, is the following true? \begin{equation}
                \Vert f^*(x) - f^*(y)\rVert_2 < \Vert f^*(x) - f^*(z)\rVert_2.
            \end{equation}
            
            \item \textbf{(0.5pt)} For random  $x \sim p_x$, can the color of $x$ be decided by $f^*(x)$ alone (\hint{ For the technical reader, we are asking if this can be decided with probability $1$. Same for the following questions.)}? 
            
            \item \textbf{(0.5pt)} For  random $x \sim p_x$, can the color of $x$ be decided by the colors of $f^*(x)$'s nearest neighbors (in training set) alone?
            
            (To find such nearest neighbors, we compute representations of all training samples, and take the samples whose representations are closest to $f^*(x)$.)
            
            \item \textbf{(0.5pt)} For  random $x \sim p_x$, can the background color of $x$ be decided by $f^*(x)$ alone? 
            
        \end{enumerate}
        
        \item \textbf{(1pt; $\mathcal{L}_\mathsf{AE}$ Implementation)} In colab, understand the parts from beginning to \hltexttt{train\_autoencoder}. 
        
        \textbf{Deliverables:} Implement the \hltexttt{FIXME} in \hltexttt{train\_autoencoder} that computes MSE reconstruction loss within 5 lines of code. 
        
        Attach your code to write-up.
                
        \item \label{rep:ae-nn-viz} \textbf{(1pt; Visualize Encoders via Nearest Neighbors)}
        In colab, we provide a function \hltexttt{nn\_visualize} that visualizes the nearest neighbors of training samples and/or unseen validation samples, where the nearest neighbors are selected based on embedding distances of a learned encoder.
        
        \textbf{Deliverables:} Run the provided code cells to \begin{enumerate}
            \item Train an encoder-decoder pair using Autoencoder objective;
            \item Visualize the trained encoder via the nearest neighbors for training samples and validation samples.
        \end{enumerate} 
        
        Attach your two plots to write-up.
                
        \item \textbf{(1pt; Understand Nearest Neighbors)} Based on the nearest neighbors plots you obtain above, answer the following questions in write-up in 2-3 sentences: \begin{enumerate}
            \item Is the representation meaningful for training images? Validation images? In particular, are similar samples grouped together?            
            
            \item The Autoencoder objective alone (on two generic functions) doesn't enforce any grouping or smoothness of the representation space. If you observe some clusters of similar images, why do you think it happens? If you do not, write down some ideas of encouraging grouping/smoothness. Your answer shouldn't be more than about 8 lines. 
                        
        \end{enumerate}
        
    \end{enumerate}
    
    
    \item \textbf{(6pt; Contrastive Learning)}
    In this part, we consider the more standard contrastive loss (not the one we used above for theoretical analysis).
    
    For encoder $f\colon \mathsf{data} \rightarrow \mathbb{S}^{d-1}$ and a temperature hyperparameter $\tau > 0$, the contrastive loss we use is:
    \begin{equation}
        \mathcal{L}_\mathsf{contr}(f; \tau) = \mathbb{E}_{\substack{(x, x^+) \sim p_{\mathsf{positive}} \\(y^{\_}_1, \dots, y^{\_}_K) \overset{\mathsf{i.i.d.}}{\sim} \mathcal{D}_\mathsf{train}}}\bigg[ %
            - \log \frac{e^{f(x)^T f(x^+) / \tau}}{e^{f(x)^T f(x^+) / \tau} + \sum_{k=1}^K e^{f(x)^T f(y^-_k) / \tau}}%
        \bigg],
    \end{equation}
    
    where $p_\mathsf{positive}$ is the distribution of positive pair. Here the pairs are \begin{align}
        (x, x^+) \tag{positive} \\
        (x, y^-_k)\mathrlap{ \qquad\quad k \in \{1, 2, \dots, K\}.} \tag{negative}
    \end{align}
    
    Notable things about implementation:\begin{itemize}
        \item Encoder $f$ is required to output $l_2$-normalized feature vectors. 
        \item $\tau$ is a (fixed) temperature hyperparameter, often selected as $\tau \in [0.05, 0.3]$. We use $\tau=0.07$ in code.
        \item As common in image contrastive learning, the positive pairs $p_\mathsf{positive}$ are defined as two random transformed versions of the \textbf{same} underlying data sample, where the negative samples are random transforms of \textbf{different} data samples.
        \item To efficiently sample features of negative pairs, we do not independently sample negatives for each positive pair. Instead, we only fetch a batch of positive pairs $\mathbf{x}$ (of $b$ samples) and $\mathbf{x}^+$ (of $b$ samples), where for each $\mathbf{x}[i]$, \begin{itemize}
            \item $(\mathbf{x}[i], \mathbf{x}^+[i])$ forms the positive pair
            
                (i.e., two random augmentations of the same underlying image). 
            \item For $j \neq i$, $(\mathbf{x}[i], \mathbf{x}^+[j])$ form the $(b-1)$ negative pairs
        \end{itemize}
      
        This means that, for encoded latents \begin{equation}
            f(\mathbf{x}) \in \mathbb{R}^{b \times d} \text{ and } f(\mathbf{x}^+) \in \mathbb{R}^{b \times d},
        \end{equation}
        the logits (i.e., pairwise dot products) can be efficiently computed using a single matrix multiplication.
        
        This is a common trick.
    \end{itemize}
        
    \begin{enumerate}
        \item \textbf{(1pt; $\mathcal{L}_\mathsf{contrastive}$ Implementation)} \textbf{Deliverables:} In the colab, follow the implementation notes above and implement the \hltexttt{FIXME} in \hltexttt{train\_contrastive} that computes contrastive loss within 5 lines of code. 
        
        Attach your code to write-up.
                
        \item \label{rep:contr-aug-inv} \textbf{(3pt; Augmentation and Invariances)} Recall from lecture and our analysis in Question~\ref{sim:align-margin-unif} that the positive pair distributions define what the learned representation should be invariant to, while sensitive to variances in all other factors (\ie, other factors are preserved in the representations). {\color{orange} ($\star$)}
        
        Here, our positive pairs are generated by random augmentations. We consider three sets of  augmentations, visualized below. 
        
        In write-up, for each set of augmentations, answer whether the learned representation would be \textit{invariant} or \textit{sensitive} to \textbf{shape}, \textbf{location} and \textbf{color}, e.g. ``Invariant to shape, sensitive to color.'' Base your answer on the {\color{orange} ($\star$)} analysis results (not on empirically trained encoders).
        \begin{enumerate}
            \item \textbf{(1pt)} \textbf{Samples from \ul{Augmentation A}:}
            
            \includegraphics[width=\linewidth]{hw4_dataset_viz/aug_A_04.pdf}
            \includegraphics[width=\linewidth]{hw4_dataset_viz/aug_A_05.pdf}
            \includegraphics[width=\linewidth]{hw4_dataset_viz/aug_A_12.pdf}
            
            (Each row shows $14$ random augmentations of the same underlying image.)
            
            Will contrastive encoder learned \wrt \textbf{Augmentation A} be invariant or sensitive to \textbf{shape}? What about \textbf{location}? What about \textbf{color}?
                        
            \item \textbf{(1pt)} \textbf{Samples from \ul{Augmentation B}:}
            
            \includegraphics[width=\linewidth]{hw4_dataset_viz/aug_B_04.pdf}
            \includegraphics[width=\linewidth]{hw4_dataset_viz/aug_B_05.pdf}
            \includegraphics[width=\linewidth]{hw4_dataset_viz/aug_B_12.pdf}
            
            (Each row shows $14$ random augmentations of the same underlying image.)
            
            Will contrastive encoder learned \wrt \textbf{Augmentation B} be invariant or sensitive to \textbf{shape}? What about \textbf{location}? What about \textbf{color}?
                        
            \item \textbf{(1pt)} \textbf{Samples from \ul{Augmentation C}:}
            
            \includegraphics[width=\linewidth]{hw4_dataset_viz/aug_C_04.pdf}
            \includegraphics[width=\linewidth]{hw4_dataset_viz/aug_C_05.pdf}
            \includegraphics[width=\linewidth]{hw4_dataset_viz/aug_C_12.pdf}
            
            (Each row shows $14$ random augmentations of the same underlying image.)
            
            Will contrastive encoder learned \wrt \textbf{Augmentation C} be invariant or sensitive to \textbf{shape}? What about \textbf{location}? What about \textbf{color}?
            
        \end{enumerate}
        
        \item \textbf{(1pt; Visualize Encoders via Nearest Neighbors)}
        
        In colab, use provided code to \begin{itemize}
            \item Train \textbf{three} contrastive encoders, one for each set of augmentations
            \item Visualize the learned encoder via nearest neighbors.
        \end{itemize} 
        
        Attach the three nearest neighbor plots ($3$ encoders on only val set) to write-up.
                
        \item \textbf{(1pt; Understand Nearest Neighbors)}
        
        In write-up, answer the following questions in 2-3 sentences: \begin{enumerate}
            
            \item For \textbf{Augmentation B}, are the learned encoder sensitive to \textbf{shape}? To \textbf{color}? To \textbf{location}?
            
            \item Are the nearest neighbor plots consistent with your answer in Question~\ref{rep:contr-aug-inv}?
                        
            \item If they are not in some cases, think about which factors are easier/harder for the encoder network to learn, and write-down some reasons on why this discrepancy happens. 
            
        \end{enumerate}
    \end{enumerate}
\end{enumerate}

\bibliographystyle{plainnat}
\bibliography{reference}

\end{document}