\documentclass[11pt,addpoints,answers]{exam}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Commands for customizing the assignment %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\newcommand{\courseNum}{10-423/623/723}
\newcommand{\courseName}{Generative AI}
\newcommand{\courseSem}{Spring 2024}
\newcommand{\courseUrl}{\url{http://423.mlcourse.org}}
\newcommand{\hwNum}{Homework 4}
\newcommand{\hwTopic}{Multi-Modal Foundation Models}
\newcommand{\hwName}{\hwNum: \hwTopic}
\newcommand{\outDate}{October 23, 2025}
\newcommand{\dueDate}{November 3, 2025}
\newcommand{\taNames}{Maxwell, Cassandra}
\newcommand{\homeworktype}{\string written+prog}
\newcommand{\autograder}{\string yes}
\newcommand{\overleafUrl}{}

%% To HIDE SOLUTIONS (to post at the website for students), set this value to 0: \def\issoln{0}
\providecommand{\issoln}{0}
%\providecommand{\issoln}{1}

%-----------------------------------------------------------------------------
% PACKAGES AND OTHER DOCUMENT CONFIGURATIONS
%-----------------------------------------------------------------------------

\usepackage[margin=1in]{geometry}
\usepackage{amsmath, amsfonts}
\usepackage{enumerate}
\usepackage{graphicx}
\usepackage{titling}
\usepackage{url}
\usepackage{xfrac}
\usepackage{natbib}
\usepackage{amssymb}
\usepackage{amsthm}
\usepackage{paralist}
\usepackage{epstopdf}
\usepackage{tabularx}
\usepackage{longtable}
\usepackage{multirow}
\usepackage{multicol}
\usepackage[colorlinks=true,urlcolor=blue]{hyperref}
\usepackage{algorithm}
\usepackage{algorithmicx}
\usepackage[noend]{algpseudocode}
\usepackage{float}
\usepackage{enumerate}
\usepackage{array}
\usepackage{environ}
\usepackage{times}
\usepackage{textcomp}
\usepackage{caption}
\usepackage{parskip} % For NIPS style paragraphs.
\usepackage[compact]{titlesec} % Less whitespace around titles
\usepackage[inline]{enumitem} % For inline enumerate* and itemize*
\usepackage{datetime}
\usepackage{comment}
% \usepackage{minted}
\usepackage{lastpage}
\usepackage{color}
% \usepackage{xcolor}
\usepackage[dvipsnames]{xcolor}
\usepackage[final]{listings}
\usepackage{tikz}
\usetikzlibrary{shapes,decorations}
\usepackage{framed}
\usepackage{booktabs}
\usepackage{cprotect}
\usepackage{verbatimbox}
\usepackage{multicol}
\usepackage{hyperref}
\usepackage{subcaption}
\usepackage{mathtools} % For drcases
\usepackage{cancel}
\usepackage[many]{tcolorbox}
\tcbuselibrary{listings} % For listings within tcolorbox
\usepackage{soul}
\usepackage[bottom]{footmisc}
\usepackage{bm}
\usepackage{wasysym}
\usepackage{lipsum}

\usepackage{tikz}
\usetikzlibrary{arrows}
\usetikzlibrary{arrows.meta}
\usetikzlibrary{shapes.geometric}
\usetikzlibrary{positioning, arrows, automata, calc}
\usepackage{transparent}
\usepackage{tikz-cd}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Formatting for \CorrectChoice of "exam" %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\CorrectChoiceEmphasis{}
\checkedchar{\blackcircle}

\newenvironment{checkboxessquare}{
    \begingroup
    \checkboxchar{$\Box$} \checkedchar{$\blacksquare$} % change checkbox style locally
    \begin{checkboxes}
    }{
    \end{checkboxes}
    \endgroup
    }

    
\newenvironment{oneparcheckboxessquare}{
    \begingroup
    \checkboxchar{$\Box$} \checkedchar{$\blacksquare$} % change checkbox style locally
    \begin{oneparcheckboxes}
    }{
    \end{oneparcheckboxes}
    \endgroup
    }

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Better numbering                        %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% \numberwithin{equation}{section} % Number equations within sections (i.e. 1.1, 1.2, 2.1, 2.2 instead of 1, 2, 3, 4)
% \numberwithin{figure}{section} % Number figures within sections (i.e. 1.1, 1.2, 2.1, 2.2 instead of 1, 2, 3, 4)
% \numberwithin{table}{section} % Number tables within sections (i.e. 1.1, 1.2, 2.1, 2.2 instead of 1, 2, 3, 4)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Custom commands                        %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newcommand{\blackcircle}{\tikz\draw[black,fill=black] (0,0) circle (1ex);}
\renewcommand{\circle}{\tikz\draw[black] (0,0) circle (1ex);}

\newcommand{\solo}{ \textcolor{orange}{[SOLO]} }
\newcommand{\open}{ \textcolor{blue}{[OPEN]} }


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Custom commands for Math               %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newcommand{\vc}[1]{\boldsymbol{#1}}
\newcommand{\adj}[1]{\frac{\partial \ell}{\partial #1}}
\newcommand{\chain}[2]{\adj{#2} = \adj{#1}\frac{\partial #1}{\partial #2}}
\newcommand{\ntset}{test}
\newcommand{\zerov}{\mathbf{0}}
\DeclareMathOperator*{\argmin}{argmin}

% mathcal
\newcommand{\Ac}{\mathcal{A}}
\newcommand{\Bc}{\mathcal{B}}
\newcommand{\Cc}{\mathcal{C}}
\newcommand{\Dc}{\mathcal{D}}
\newcommand{\Ec}{\mathcal{E}}
\newcommand{\Fc}{\mathcal{F}}
\newcommand{\Gc}{\mathcal{G}}
\newcommand{\Hc}{\mathcal{H}}
\newcommand{\Ic}{\mathcal{I}}
\newcommand{\Jc}{\mathcal{J}}
\newcommand{\Kc}{\mathcal{K}}
\newcommand{\Lc}{\mathcal{L}}
\newcommand{\Mc}{\mathcal{M}}
\newcommand{\Nc}{\mathcal{N}}
\newcommand{\Oc}{\mathcal{O}}
\newcommand{\Pc}{\mathcal{P}}
\newcommand{\Qc}{\mathcal{Q}}
\newcommand{\Rc}{\mathcal{R}}
\newcommand{\Sc}{\mathcal{S}}
\newcommand{\Tc}{\mathcal{T}}
\newcommand{\Uc}{\mathcal{U}}
\newcommand{\Vc}{\mathcal{V}}
\newcommand{\Wc}{\mathcal{W}}
\newcommand{\Xc}{\mathcal{X}}
\newcommand{\Yc}{\mathcal{Y}}
\newcommand{\Zc}{\mathcal{Z}}

% mathbb
\newcommand{\Ab}{\mathbb{A}}
\newcommand{\Bb}{\mathbb{B}}
\newcommand{\Cb}{\mathbb{C}}
\newcommand{\Db}{\mathbb{D}}
\newcommand{\Eb}{\mathbb{E}}
\newcommand{\Fb}{\mathbb{F}}
\newcommand{\Gb}{\mathbb{G}}
\newcommand{\Hb}{\mathbb{H}}
\newcommand{\Ib}{\mathbb{I}}
\newcommand{\Jb}{\mathbb{J}}
\newcommand{\Kb}{\mathbb{K}}
\newcommand{\Lb}{\mathbb{L}}
\newcommand{\Mb}{\mathbb{M}}
\newcommand{\Nb}{\mathbb{N}}
\newcommand{\Ob}{\mathbb{O}}
\newcommand{\Pb}{\mathbb{P}}
\newcommand{\Qb}{\mathbb{Q}}
\newcommand{\Rb}{\mathbb{R}}
\newcommand{\Sb}{\mathbb{S}}
\newcommand{\Tb}{\mathbb{T}}
\newcommand{\Ub}{\mathbb{U}}
\newcommand{\Vb}{\mathbb{V}}
\newcommand{\Wb}{\mathbb{W}}
\newcommand{\Xb}{\mathbb{X}}
\newcommand{\Yb}{\mathbb{Y}}
\newcommand{\Zb}{\mathbb{Z}}

% mathbf lowercase
\newcommand{\av}{\mathbf{a}}
\newcommand{\bv}{\mathbf{b}}
\newcommand{\cv}{\mathbf{c}}
\newcommand{\dv}{\mathbf{d}}
\newcommand{\ev}{\mathbf{e}}
\newcommand{\fv}{\mathbf{f}}
\newcommand{\gv}{\mathbf{g}}
\newcommand{\hv}{\mathbf{h}}
\newcommand{\iv}{\mathbf{i}}
\newcommand{\jv}{\mathbf{j}}
\newcommand{\kv}{\mathbf{k}}
\newcommand{\lv}{\mathbf{l}}
\newcommand{\mv}{\mathbf{m}}
\newcommand{\nv}{\mathbf{n}}
\newcommand{\ov}{\mathbf{o}}
\newcommand{\pv}{\mathbf{p}}
\newcommand{\qv}{\mathbf{q}}
\newcommand{\rv}{\mathbf{r}}
\newcommand{\sv}{\mathbf{s}}
\newcommand{\tv}{\mathbf{t}}
\newcommand{\uv}{\mathbf{u}}
\newcommand{\vv}{\mathbf{v}}
\newcommand{\wv}{\mathbf{w}}
\newcommand{\xv}{\mathbf{x}}
\newcommand{\yv}{\mathbf{y}}
\newcommand{\zv}{\mathbf{z}}

% mathbf uppercase
\newcommand{\Av}{\mathbf{A}}
\newcommand{\Bv}{\mathbf{B}}
\newcommand{\Cv}{\mathbf{C}}
\newcommand{\Dv}{\mathbf{D}}
\newcommand{\Ev}{\mathbf{E}}
\newcommand{\Fv}{\mathbf{F}}
\newcommand{\Gv}{\mathbf{G}}
\newcommand{\Hv}{\mathbf{H}}
\newcommand{\Iv}{\mathbf{I}}
\newcommand{\Jv}{\mathbf{J}}
\newcommand{\Kv}{\mathbf{K}}
\newcommand{\Lv}{\mathbf{L}}
\newcommand{\Mv}{\mathbf{M}}
\newcommand{\Nv}{\mathbf{N}}
\newcommand{\Ov}{\mathbf{O}}
\newcommand{\Pv}{\mathbf{P}}
\newcommand{\Qv}{\mathbf{Q}}
\newcommand{\Rv}{\mathbf{R}}
\newcommand{\Sv}{\mathbf{S}}
\newcommand{\Tv}{\mathbf{T}}
\newcommand{\Uv}{\mathbf{U}}
\newcommand{\Vv}{\mathbf{V}}
\newcommand{\Wv}{\mathbf{W}}
\newcommand{\Xv}{\mathbf{X}}
\newcommand{\Yv}{\mathbf{Y}}
\newcommand{\Zv}{\mathbf{Z}}

% bold greek lowercase
\newcommand{\alphav     }{\boldsymbol \alpha     }
\newcommand{\betav      }{\boldsymbol \beta      }
\newcommand{\gammav     }{\boldsymbol \gamma     }
\newcommand{\deltav     }{\boldsymbol \delta     }
\newcommand{\epsilonv   }{\boldsymbol \epsilon   }
\newcommand{\varepsilonv}{\boldsymbol \varepsilon}
\newcommand{\zetav      }{\boldsymbol \zeta      }
\newcommand{\etav       }{\boldsymbol \eta       }
\newcommand{\thetav     }{\boldsymbol \theta     }
\newcommand{\varthetav  }{\boldsymbol \vartheta  }
\newcommand{\iotav      }{\boldsymbol \iota      }
\newcommand{\kappav     }{\boldsymbol \kappa     }
\newcommand{\varkappav  }{\boldsymbol \varkappa  }
\newcommand{\lambdav    }{\boldsymbol \lambda    }
\newcommand{\muv        }{\boldsymbol \mu        }
\newcommand{\nuv        }{\boldsymbol \nu        }
\newcommand{\xiv        }{\boldsymbol \xi        }
\newcommand{\omicronv   }{\boldsymbol \omicron   }
\newcommand{\piv        }{\boldsymbol \pi        }
\newcommand{\varpiv     }{\boldsymbol \varpi     }
\newcommand{\rhov       }{\boldsymbol \rho       }
\newcommand{\varrhov    }{\boldsymbol \varrho    }
\newcommand{\sigmav     }{\boldsymbol \sigma     }
\newcommand{\varsigmav  }{\boldsymbol \varsigma  }
\newcommand{\tauv       }{\boldsymbol \tau       }
\newcommand{\upsilonv   }{\boldsymbol \upsilon   }
\newcommand{\phiv       }{\boldsymbol \phi       }
\newcommand{\varphiv    }{\boldsymbol \varphi    }
\newcommand{\chiv       }{\boldsymbol \chi       }
\newcommand{\psiv       }{\boldsymbol \psi       }
\newcommand{\omegav     }{\boldsymbol \omega     }

% bold greek uppercase
\newcommand{\Gammav     }{\boldsymbol \Gamma     }
\newcommand{\Deltav     }{\boldsymbol \Delta     }
\newcommand{\Thetav     }{\boldsymbol \Theta     }
\newcommand{\Lambdav    }{\boldsymbol \Lambda    }
\newcommand{\Xiv        }{\boldsymbol \Xi        }
\newcommand{\Piv        }{\boldsymbol \Pi        }
\newcommand{\Sigmav     }{\boldsymbol \Sigma     }
\newcommand{\Upsilonv   }{\boldsymbol \Upsilon   }
\newcommand{\Phiv       }{\boldsymbol \Phi       }
\newcommand{\Psiv       }{\boldsymbol \Psi       }
\newcommand{\Omegav     }{\boldsymbol \Omega     }

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Code highlighting with listings         %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\definecolor{bluekeywords}{rgb}{0.13,0.13,1}
\definecolor{greencomments}{rgb}{0,0.5,0}
\definecolor{redstrings}{rgb}{0.9,0,0}
\definecolor{light-gray}{gray}{0.95}

\newcommand{\MYhref}[3][blue]{\href{#2}{\color{#1}{#3}}}%

\definecolor{dkgreen}{rgb}{0,0.6,0}
\definecolor{gray}{rgb}{0.5,0.5,0.5}
\definecolor{mauve}{rgb}{0.58,0,0.82}
\definecolor{antiquefuchsia}{rgb}{0.57, 0.36, 0.51}

\lstdefinelanguage{Shell}{
  keywords={tar, cd, make},
  %keywordstyle=\color{bluekeywords}\bfseries,
  alsoletter={+},
  ndkeywords={python3, python, py, javac, java, gcc, c, g++, cpp, .txt, octave, m, .tar},
  %ndkeywordstyle=\color{bluekeywords}\bfseries,
  identifierstyle=\color{black},
  sensitive=false,
  comment=[l]{//},
  morecomment=[s]{/*}{*/},
  commentstyle=\color{purple}\ttfamily,
  %stringstyle=\color{red}\ttfamily,
  morestring=[b]',
  morestring=[b]",
  backgroundcolor = \color{light-gray}
}

\lstdefinestyle{mypython}{
    language=Python,
    basicstyle=\ttfamily\small,       % font and size
    keywordstyle=\color{blue},        % keywords
    stringstyle=\color{orange},       % strings
    commentstyle=\color{green!50!black}, % comments
    showstringspaces=false,           % don't show spaces in strings
    numberstyle=\tiny\color{gray},    % line numbers
    numbers=left,                     % line numbers on the left
    stepnumber=1,                     % number every line
    breaklines=true,                  % automatic line breaking
    frame=single,                     % box around code
    tabsize=4                         % tab width
}

\lstset{columns=fixed, basicstyle=\ttfamily,
    backgroundcolor=\color{light-gray},xleftmargin=0.5cm,frame=tlbr,framesep=4pt,framerule=0pt}

\newcommand{\emptysquare}{{\LARGE $\square$}\ \ }
\newcommand{\filledsquare}{{\LARGE $\boxtimes$}\ \ }
\def \ifempty#1{\def\temp{#1} \ifx\temp\empty }


% \newcommand{\squaresolutionspace}[2][\emptysquare]{\newline #1}{#2}
\def \squaresolutionspace#1{ \ifempty{#1} \emptysquare \else #1\hspace{0.75pt}\fi}


\newcommand{\emptycircle}{{\LARGE $\fullmoon$}\ \ }
\newcommand{\filledcircle}{{\LARGE $\newmoon$}\ \ }
\def \circlesolutionspace#1{ \ifempty{#1} \emptycircle \else #1\hspace{0.75pt}\fi}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Custom box for highlights               %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Define box and box title style
\tikzstyle{mybox} = [fill=blue!10, very thick,
    rectangle, rounded corners, inner sep=1em, inner ysep=1em]

% \newcommand{\notebox}[1]{
% \begin{tikzpicture}
% \node [mybox] (box){%
%     \begin{minipage}{\textwidth}
%     #1
%     \end{minipage}
% };
% \end{tikzpicture}%
% }

\NewEnviron{notebox}{
\begin{tikzpicture}
\node [mybox] (box){
    \begin{minipage}{0.95\textwidth}
        \BODY
    \end{minipage}
};
\end{tikzpicture}
}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Commands showing / hiding solutions     %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newcommand{\solutionspace}[4]{\fbox{\begin{minipage}[t][#1][t]{#2} \textbf{#3} \solution{}{#4} \end{minipage}}}

% To HIDE SOLUTIONS, set this value to 0:
%\providecommand{\issoln}{0}
%\providecommand{\issoln}{1}

\ifthenelse{\equal{\issoln}{1}}{

% SOLUTION environment
\newenvironment{soln}{\leavevmode\color{red}\ignorespaces }{}

% QUESTION AUTHORS environment
\newenvironment{qauthor}{\leavevmode\color{blue}\ignorespaces }{}

% Question tester comment environment
\newenvironment{qtester}{\leavevmode\color{green}\ignorespaces}{}

% Question learning objective comment environment
\newenvironment{qlearningobjective}{\leavevmode\color{green}\ignorespaces}{}

}{ % ELSE

  \NewEnviron{soln}{}
  \NewEnviron{qauthor}{}
  \NewEnviron{qtester}{}
  \NewEnviron{qlearningobjective}{}

}

%% To HIDE TAGS set this value to 0:
\def\showtags{0}
%%%%%%%%%%%%%%%%
\ifcsname showtags\endcsname \else \def\showtags{1} \fi
% Default to an empty tags environ.
\NewEnviron{tags}{}{}
\if\showtags 1
% Otherwise, include solutions as below.
\RenewEnviron{tags}{
    \fbox{
    \leavevmode\color{blue}\ignorespaces
    \textbf{TAGS:} \texttt{\url{\BODY}}
    }
    \vspace{-.5em}
}{}
\fi

\newtcolorbox[]{answer_box}[1][]
{
    % breakable,
    fit,
    enhanced,
    % nobeforeafter,
    colback=white,
    title=Your Answer,
    sidebyside align=top,
    box align=top,
    #1
}

% This must be specified with option brackets \begin{answerboxcode}[] \end{answerboxcode}
% For some reason, this would break if we used the name answer_box_code
\newtcblisting{answerboxcode}[1][]
{
    listing only,
    listing options={
        language=Python,
        showstringspaces=false,     % Don't show spaces in strings
        tabsize=4,                  % Set tab size to 4 spaces
        breaklines=true,            % Break long lines
        numbers=left,               % Line numbers on the left
    },
    height=6cm,
    width=15cm,
    fit,
    enhanced,
    colback=white,
    title=Your Answer,
    sidebyside align=top,
    box align=top,
    #1
}

%\pagestyle{fancyplain}
\lhead{\hwName{}
}
\rhead{\courseNum}
\cfoot{\thepage{} of \numpages{}}

\title{\textsc{\hwNum}
\\
\textsc{\hwTopic}
\thanks{Compiled on \today{} at \currenttime{}}\\
\vspace{1em}
} % Title


\author{\textsc{\large \courseNum{} \courseName{}}\\
\courseUrl
\vspace{1em}\\
\ifdefempty{\outDate}{}{  OUT: \outDate \\ }
\ifdefempty{\dueDate}{}{  DUE: \dueDate \\ }
\ifdefempty{\taNames}{}{  TAs: \taNames{} }
}

\date{}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Useful commands for typesetting the questions %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% This command will allow long \lstinline{} text to wrap automatically.
\sloppy

%%%%%%%%%%%%%%%%%%%%%%%%%%
% Document configuration %
%%%%%%%%%%%%%%%%%%%%%%%%%%

% Don't display a date in the title and remove the white space
\predate{}
\postdate{}
\date{}

% Don't display an author and remove the white space
% \preauthor{}
% \postauthor{}

% Question type commands
\newcommand{\sall}{\textbf{Select all that apply: }}
\newcommand{\sone}{\textbf{Select one: }}
\newcommand{\tf}{\textbf{True or False: }}




% Changes to examdoc
\qformat{\textbf{{\Large \thequestion \; \; \thequestiontitle \ (\totalpoints \ points)}} \hfill}
\renewcommand{\thequestion}{\arabic{question}}
\renewcommand{\questionlabel}{\thequestion.}

\renewcommand{\thepartno}{\thequestion.\arabic{partno}}
%\renewcommand{\partlabel}{\thequestion.\thepartno.}
\renewcommand{\partlabel}{\thepartno.}

% not working: \renewcommand{\subpartlabel}{(\thequestion.\thepartno.\thesubpart)}
% Commented after adding \question.\thepartno.
%\renewcommand{\partshook}{\setlength{\leftmargin}{0pt}}

\renewcommand{\thesubpart}{\thepartno.\alph{subpart}}
\renewcommand{\subpartlabel}{\thesubpart.}

\renewcommand{\thesubsubpart}{\thesubpart.\roman{subsubpart}}
\renewcommand{\subsubpartlabel}{\thesubsubpart.}

% copied from stack overflow, as all good things are
\newcommand\invisiblesection[1]{%
  \refstepcounter{section}%
  \addcontentsline{toc}{section}{\protect\numberline{\thesection}#1}%
  \sectionmark{#1}}

% quite possibly the worst workaround i have made for this class
\newcommand{\sectionquestion}[1]{
\titledquestion{#1}
\invisiblesection{#1}
~\vspace{-1em}
}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% New Environment for Pseudocode          %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Python style for highlighting
\DeclareFixedFont{\ttb}{T1}{txtt}{bx}{n}{12} % for bold
\DeclareFixedFont{\ttm}{T1}{txtt}{m}{n}{12}  % for normal

\definecolor{deepblue}{rgb}{0,0,0.5}
\definecolor{deepred}{rgb}{0.6,0,0}
\definecolor{deepgreen}{rgb}{0,0.5,0}

\newcommand\pythonstyle{\lstset{
language=Python,
basicstyle=\ttm,
morekeywords={self},              % Add keywords here
keywordstyle=\ttb\color{deepblue},
emph={MyClass,__init__},          % Custom highlighting
emphstyle=\ttb\color{deepred},    % Custom highlighting style
stringstyle=\color{deepgreen},
frame=tb,                         % Any extra options here
showstringspaces=false
}}


% Python environment
\lstnewenvironment{your_code_solution}[1][]
{
\pythonstyle
\lstset{#1}
}
{}

\begin{document}
 
\maketitle 

\newcommand \maxsubs {10 }
\section*{Instructions}
\begin{itemize}

\item \textbf{Collaboration Policy}: Please read the collaboration policy in the syllabus.
\item\textbf{Late Submission Policy:} See the late submission policy in the syllabus.
\item\textbf{Submitting your work:} You will use Gradescope to submit
  answers to all questions\ifthenelse{\equal{\homeworktype}{\string written}}{}{ and code}. 

\begin{itemize}
    
    % IF NOT USING TEMPLATE: 
    % \item \textbf{Written:} You will submit your completed homework as a PDF to Gradescope. For each problem, please clearly indicate the question number (e.g. 3.2). Submissions can be handwritten, but must be clearly legible; otherwise, you will not be awarded marks.   Alternatively, submissions can be written in \LaTeX{}. You may use the \LaTeX{} source of this assignment (included in the handout .zip) as your starting point. For multiple choice / select all questions, simply write the letter(s) (e.g. A, B, C) corresponding to your chosen answer.
    % IF USING TEMPLATE: 
    \item \textbf{Written:} You will submit your completed homework as a PDF to Gradescope. Please use the provided template. Submissions can be handwritten, but must be clearly legible; otherwise, you will not be awarded marks. Alternatively, submissions can be written in \LaTeX{}. Each answer should be within the box provided. 
    %If you do not follow the template or your submission is misaligned, your assignment may not be graded correctly by our AI assisted grader. 
    If you do not follow the template, your assignment may not be graded correctly by our AI assisted grader and there will be a \textbf{\textcolor{red}{2\% penalty}} (e.g., if the homework is out of 100 points, 2 points will be deducted from your final score).
    
    \ifdefempty{\overleafUrl}{}{
    \item \textbf{\LaTeX{} Source:} \overleafUrl
    }

    \ifthenelse{\equal{\homeworktype}{\string written}}{}{
    \item \textbf{Programming:} You will submit your code for programming questions to Gradescope. \ifthenelse{\equal{\autograder}{\string yes}}{}{ There is no autograder. }
    We will examine your code by hand and may award marks for its submission.
    }{}
   
  \end{itemize}
  
\ifthenelse{\equal{\homeworktype}{\string written}}{}{\item\textbf{Materials:} The data that you will need in order to complete this assignment is posted along with the writeup and template on the course website.}

\end{itemize}

{\small
\begin{center}
    \pointtable[v][questions]
\end{center}
}\clearpage

%\input{../shared/instructions_for_specific_problem_types.tex}
%\clearpage
\begin{questions}

\sectionquestion{\LaTeX{} Template Alignment}
\begin{parts}
    \part[0] \sone Did you use \LaTeX{} for the entire written portion of this homework?
    
    \begin{checkboxes}
        % YOUR ANSWER
        % Change \choice to \CorrectChoice for the appropriate selection/selections 
        \choice Yes 
        \choice No
    \end{checkboxes}

    \part[0] \sone I have ensured that my final submission is aligned with the original template given to me in the handout file and that I haven't deleted or resized any items or made any other modifications which will result in a misaligned template. I understand that incorrectly responding yes to this question will result in a penalty equivalent to 2\% of the points on this assignment.\\
    \textbf{Note:} Failing to answer this question will not exempt you from the 2\% misalignment penalty.
    
    \begin{checkboxes}
        % YOUR ANSWER
        % Change \choice to \CorrectChoice for the appropriate selection/selections 
        \choice Yes 
    \end{checkboxes}
\end{parts}\clearpage

\sectionquestion{Latent Diffusion Model (LDM)}

\begin{parts}
\part[1] \textbf{Math:} A latent diffusion model uses a compression factor of $f=8$ to downsample. If an input image is $512\times512$ pixels, what are the dimensions of the latent representation of this image?
    \begin{answer_box}[title=,height=1.5cm, width=7.5cm]
    \end{answer_box}
    

\part[2] \textbf{Short answer:} Why does a latent diffusion model run diffusion in a latent space instead of pixel space?
    \begin{answer_box}[title=,height=2.5cm, width=15cm]
    \end{answer_box}

\part \textbf{Short answer:} Standard cross-attention for a diffusion-based text-to-image model defines the queries $\Qv$ as a function of the pixels (or latent space) $\Yv \in \Rb^{m \times d_y}$, and the keys $\Kv$ and values $\Vv$ as a function of the text encoder output $\Xv \in \Rb^{n \times d_x}$.
\begin{align*}
    &&\Qv &= \Yv \Wv_q,
    &\Kv &= \Xv \Wv_k, 
    &\Vv &= \Xv \Wv_v
\end{align*}
(where $\Wv_q \in \Rb^{d_y \times d}$
and $\Wv_k, \Wv_v \in \Rb^{d_x \times d}$)
and then applies standard attention:
\begin{align*}
    \text{Attention}(\Qv, \Kv, \Vv) = \text{softmax}(\Qv \Kv^T / \sqrt{d}) \Vv
\end{align*}

Now, suppose you instead defined a new formulation where the values are a function of the pixels (or latent space): $\Vv = \Yv \Wv_v$ where $\Wv_v \in \Rb^{d_y \times d}$.

\begin{subparts}

\subpart[2] What goes wrong mathematically in the new formulation?
    \begin{answer_box}[title=,height=3cm, width=15cm]
    \end{answer_box}
    
\subpart[2] Intuitively, why doesn't the new formulation make sense? Briefly begin with an explanation of what the original formulation of cross-attention is trying to accomplish for a single query vector, and why this new formulation fails to accomplish that.
    \begin{answer_box}[title=,height=4cm, width=15cm]
    \end{answer_box}
    
\end{subparts}




\end{parts}

\clearpage

\sectionquestion{VQ-VAEs}

\begin{parts}
\part The objective function for a VQ-VAE contains two terms in addition to the reconstruction loss term: 
\begin{itemize}
    \item The vector-quantization loss, $\lVert\textrm{sg}\left[z_e(x)\right]-e\rVert_2^2$ and
    \item The ``commitment'' loss, $\beta \lVert z_e(x)-\textrm{sg}\left[e\right]\rVert_2^2$
\end{itemize}
where sg is the stopgradient operator and $e$ is the latent embedding vector closest to the output of the encoder $z_e(x)$.\footnote{For the purposes of this question, we are still assuming that the encoder outputs a single vector; in practice, a true VQ-VAE encoder would output multiple vectors and these terms in the objective function would be sums over these vectors.}

In this question, you will examine the impact of these terms and the stopgradient operator. Let 
\begin{align}
    \mathcal{L} &=\lVert\textrm{sg}\left[z_e(x)\right]-e\rVert_2^2+\beta \lVert z_e(x)-\textrm{sg}\left[e\right]\rVert_2^2 \\
    \tilde{\mathcal{L}} &= \lVert z_e(x)-e\rVert_2^2+\beta \lVert z_e(x)-e\rVert_2^2
\end{align}

\begin{subparts}
    \subpart[1] \textbf{Math:} What is the gradient of $\mathcal{L}$ with respect to $z_e(x)$?
    \begin{answer_box}[title=,height=1.5cm, width=7.5cm]
    \end{answer_box}
    
    \subpart[1] \textbf{Math:} What is the gradient of $\mathcal{L}$ with respect to $e$?
    \begin{answer_box}[title=,height=1.5cm, width=7.5cm]
    \end{answer_box}
    
    \subpart[1] \textbf{Math:} What is the gradient of $\tilde{\mathcal{L}}$ with respect to $z_e(x)$?
    \begin{answer_box}[title=,height=1.5cm, width=7.5cm]
    \end{answer_box}
    
    \subpart[1] \textbf{Math:} What is the gradient of $\tilde{\mathcal{L}}$ with respect to $e$?
    \begin{answer_box}[title=,height=1.5cm, width=7.5cm]
    \end{answer_box}

    \subpart[2] \textbf{Short answer:} Given your findings from the previous parts, how would you describe the impact of the stopgradient operator on the optimization of these two terms? How do the gradients with and without the stopgradient operator differ?
    \begin{answer_box}[title=,height=2cm, width=15cm]
    \end{answer_box}
\end{subparts}

\clearpage
\part[2] In a VQ-VAE, the codebook embeddings $e$ need to be learned during training. Recall that the VQ-VAE objective includes the vector-quantization loss $\lVert\textrm{sg}\left[z_e(x)\right]-e\rVert_2^2$, where $z_e(x)$ is the encoder output, $e$ is the nearest codebook embedding, and $\textrm{sg}\left[ \cdot \right]$ is the stop gradient operator. Using your answer from 3.1, describe in words what the codebook embeddings are being optimized to do.

    \begin{answer_box}[title=,height=2.5cm, width=15cm]
    \end{answer_box}
    
    

\end{parts}

\clearpage
\sectionquestion{CLIP}

\begin{parts}
\part[4] \textbf{Short answer:} In \href{https://arxiv.org/pdf/2103.00020}{Radford et al. (2021)}, the original CLIP paper, they pre-trained the image and text encoders using a \emph{symmetric cross-entropy loss}. Formally, given a mini-batch of $N$ (image, caption) pairs, let $e_i^{(n)}$ be the image embedding of the $n^{\textrm{th}}$ image normalized to have unit norm and let $e_t^{(n)}$ be the text embedding of the $n^{\textrm{th}}$ caption, also normalized to have unit norm. 

The symmetric cross-entropy loss consists of two terms for each (image, caption) pair, a term that compares the image embedding to all $N$ caption embeddings and a term that compares the caption embedding to all $N$ image embeddings. Formally, let
\begin{align}
    \ell_{i\to t}^{(n)} &= -\log\left(\frac{\exp\left(e_i^{(n)}\cdot e_t^{(n)}\right)}{\sum_{m=1}^N \exp\left(e_i^{(n)}\cdot e_t^{(m)}\right)}\right) \\
    \ell_{t\to i}^{(n)} &= -\log\left(\frac{\exp\left(e_i^{(n)}\cdot e_t^{(n)}\right)}{\sum_{m=1}^N \exp\left(e_i^{(m)}\cdot e_t^{(n)}\right)}\right) 
\end{align}
where $\cdot$ is the dot-product operation between two (same-length) vectors. The CLIP objective can then be expressed as 
\begin{equation}
    \ell = \frac{1}{N}\sum_{n=1}^N\frac{\ell_{i\to t}^{(n)}+\ell_{t\to i}^{(n)}}{2}
\end{equation}

Show that minimizing this objective function is equivalent to jointly maximizing the cosine similarity of an image embedding and its corresponding caption embedding while minimizing the cosine similarity of all other pairs of image and caption embeddings. 

\textbf{Hint:} \textit{It may be helpful to start with $\min(\ell)$, and follow through with the mathematical logic until you show that it evaluates to $\max_{n} \{e_i^{(n)} \cdot e_t^{(n)}\} \text{ and } \min_{n \neq m} \{e_i^{(n)} \cdot e_t^{(m)}, e_i^{(m)} \cdot e_t^{(n)}\}$}.
\begin{answer_box}[title=,height=6cm, width=15cm]
\end{answer_box}
\end{parts}

\clearpage

\sectionquestion{VLMs}

This section will familiarize you with the working of VLMs. For simplicity, we will be focusing on the PaliGemma2 Model. We encourage you to read this \href{https://arxiv.org/abs/2412.03555}{Research Paper} and this \href{https://huggingface.co/blog/paligemma2}{Article} to better understand the details of the model. 

\begin{parts}


\part[1] \textbf{Select one:} Which component converts images into visual tokens in PaliGemma 2?
    \begin{checkboxes}
    \choice Gemma 2 language model
    \choice SigLIP vision encoder
    \choice Linear projection layer
    \choice Autoregressive decoder
    \end{checkboxes}

\part[1] \textbf{Select one:} What is the primary purpose of the projection layer in multimodal models like PaliGemma2?
    \begin{checkboxes}
    \choice To downsample the image size for faster processing
    \choice To align vision encoder outputs with language model's embedding space
    \choice To generate more high-fidelity images from text descriptions
    \choice To filter out irrelevant visual information before training attention layers
    \end{checkboxes}

\part[1] \textbf{Select one:} Why does PaliGemma2 \textit{primarily} use a pre-trained vision encoder like SigLIP instead of training a vision component from scratch?
    \begin{checkboxes}
    \choice To save on computational resources only
    \choice Because training vision models is impossible within this framework
    \choice To leverage the visual understanding capabilities  present in pre-trained vision models
    \choice Because processing raw images directly with neural networks leads to worse output
    \end{checkboxes}

\part[1] \textbf{Select one:} What advantage does using a pre-trained language model like Gemma provide in the PaliGemma2 architecture?
    \begin{checkboxes}
    \choice It eliminates the need for any image processing components
    \choice It brings language understanding capabilities to visual reasoning tasks
    \choice It makes the system completely language-independent
    \choice It automatically generates training data for the vision encoder
    \end{checkboxes}

\clearpage
\part[2] \textbf{Short Answer:} With an image resolution of 224×224 pixels and a patch size of 14×14 pixels, Pali-Gemma 2 generates 256 visual tokens. How would the number of tokens change if we increased the resolution to 448×448 pixels while maintaining the same patch size of 14×14 pixels? Why is this change important for the model's performance?
    \begin{answer_box}[title=,height=3cm, width=15cm]
    \end{answer_box}

\part[2] \textbf{Short Answer:} VLMs are great for tasks that they have been trained on, but still can be tricked by adversarial examples. For this question, first take a normal image and ask a proprietary VLM like GPT or Gemini to densely caption it (i.e., describe what's happening in the image in detail), and paste the dense caption. Next, find/create an adversarial image and question pair that a VLM gets completely wrong. Attach a screenshot of this exchange, like the example below:
\begin{figure}[H]
        \centering
        \includegraphics[width=.6\linewidth]{fig/VLM_counterfactual.png}       
        \label{fig:VLMcounterfactual}
    \end{figure}
  \begin{answer_box}[title=Dense Caption,height=2cm, width=15cm]
    \end{answer_box}
    \begin{answer_box}[title=Screenshot,height=6cm, width=15cm]
    \end{answer_box}


\end{parts}



\clearpage

\clearpage
\sectionquestion{Programming: Text-to-Image Generation}

\uplevel{\subsection*{Introduction}} 

Text-to-image generation has revolutionized how we create visual content, with models like DALL-E, Stable Diffusion, and Midjourney demonstrating remarkable capabilities. These large-scale models require extensive computational resources and massive datasets. Such models have moved us closer to the objective of being able to combine the knowledge of an LLM with the visual quality of a image generation model for a chat-like generation experience. In this assignment, you'll explore the fundamental principles of LLM-conditioned image generation by combining a pretrained LLM and class-conditional diffusion model into a single text-to-image generation model.

The core challenge in text-to-image generation is bridging the semantic gap between language and vision. While diffusion models excel at generating high-quality images and language models capture rich semantic information, connecting these modalities efficiently remains non-trivial. In this assignment, you will train a \textbf{Q-Former} (Query Former) architecture proposed by \href{https://arxiv.org/pdf/2107.14795}{Perciever IO}, and used to connect text and image features in \href{https://arxiv.org/abs/2301.12597}{BLIP-2 (Li et al., 2023)} and \href{https://arxiv.org/pdf/2504.06256}{MetaQueries}. The Q-Former uses a fixed number of learnable query tokens to extract relevant visual features from arbitrary length language model representations.

Your will:
\begin{itemize}
    \item Train a Q-Former to use language model text embeddings from GPT-2 as conditioning for a pretrained diffusion model with a Diffusion Transformer (DiT) backbone, while keeping both GPT-2 and the DiT model weights frozen.
    \item Generate 32×32 images from simple text prompts like "fluffy cat" or "red airplane".
    \item Examine cross-attention maps to determine what information the Q-Former is extracting from the pretrained LLM.
\end{itemize}
\begin{notebox}
\textbf{Note:} CIFAR-10 images are only 32×32 pixels, so generated images may appear blurry or lack fine details. This is expected due to the low resolution, not a sign of poor model performance. We highly recommend that you take a look at CIFAR-10 for reference.
\end{notebox}

\uplevel{\subsection*{Background}}

\textbf{The Text-to-Image Pipeline:}

Our goal is to generate images from text descriptions. We have three components:

\begin{enumerate}
    \item \textbf{GPT-2 (Frozen)}: A pretrained language model that converts text into semantic feature vectors
    \item \textbf{DiT Diffusion Model (Frozen)}: A pretrained class-conditional diffusion transformer trained on the \href{https://www.cs.toronto.edu/~kriz/cifar.html}{CIFAR dataset}
    \item \textbf{Q-Former (Trained)}: A lightweight adapter that translates language features into visual conditioning, allowing us to provide precise captions for cifar-like images, and extrapolate beyond class-based generation
\end{enumerate}

\definecolor{antiquefuchsia}{rgb}{0.57, 0.36, 0.51}
\begin{figure}[H]
    \centering
    \includegraphics[width=\linewidth]{fig/Qformer_Figure.png}
    \caption{Our Q-Former architecture for text-to-image generation.}
    \label{fig:qformer}
\end{figure}

First, descriptive text is tokenized and passed through GPT-2 to produce \textcolor{orange}{\textbf{per-token text features}}
 (bottom-left, white boxes up to orange layer). Next, we initialize \textcolor{antiquefuchsia}{\textbf{learnable query embeddings}}, and transform them with transformer blocks that perform self-attention, cross-attention with the \textcolor{orange}{\textbf{per-token text features}}, and a per-token feed-forward network (FFN) (top-left, orange layer to purple blocks). Then, a DiT projector outputs a set of \textcolor{antiquefuchsia}{\textbf{class embedding features}} (middle, purple block). During the forward pass of the Diffusion Transformer (DiT), the class embedding vector for the current image class is replaced with the \textcolor{antiquefuchsia}{\textbf{per-layer class embedding features}} at every layer (from purple middle and top-right noisy image to bottom-right predicted noise). 

 
 \textbf{During training, only the \textcolor{antiquefuchsia}{purple layers} are trained, while all other weights are frozen.}

\textbf{Why Q-Former?}

Why don't we use GPT-2 features directly as conditioning for the diffusion model? Here are some challenges:

\begin{itemize}
    \item GPT-2 outputs a variable length sequence of 768-dimensional vectors, but DiT expects a single 768-dimensional vector for conditioning.
    \item GPT-2 features live in ``language space" while DiT needs ``visual conditioning space," and we would not expect these representations to be aligned in any sense.
    %\item GPT-2 produces variable-length sequences, but DiT expects fixed-size conditioning
\end{itemize}

The Q-Former solves this by learning to extract relevant information from language features and projecting it into the format DiT expects.

\textbf{Diffusion Models for Image Generation:}

Denoising Diffusion Probabilistic Models (DDPMs) \href{https://arxiv.org/abs/2006.11239}{(Ho et al., 2020)} generate images by iteratively denoising random Gaussian noise. The model learns to predict and remove noise step-by-step:

$$x_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left(x_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta(x_t, t, c)\right) + \sigma_t z$$

where $\epsilon_\theta$ is the noise prediction network (DiT), $c$ is the conditioning information, and $\alpha_t$ are noise schedule parameters.

For this assignment, we use a pretrained \href{https://arxiv.org/abs/2212.09748}{Diffusion Transformer (DiT)} trained on CIFAR-10 with class label conditioning. The sampling process is already implemented for you. To perform the conditioning, the pretrained model has a learned class embedding for each of the 10 classes, which is used to modulate each transformer layer via an Adaptive LayerNorm. This takes the class embedding as input and outputs LayerNorm scale and shift parameters:

$$x = (\text{AdALN(}\textbf{F}_{\text{class embedding}}))(x)$$


For more details, see the Adaptive LayerNorm section of the original \href{https://arxiv.org/abs/2212.09748}{Diffusion Transformer paper}, or the \lstinline{adaLN_modulation} function in \lstinline{dit.py}.

\textbf{Classifier-Free Guidance:}

To improve prompt adherence, the provided code uses classifier-free guidance \href{https://arxiv.org/abs/2207.12598}{(Ho \& Salimans, 2022)}, which amplifies the effect of conditioning:

$$\tilde{\epsilon}_\theta(x_t, t, c) = \epsilon_\theta(x_t, t, \emptyset) + w \cdot (\epsilon_\theta(x_t, t, c) - \epsilon_\theta(x_t, t, \emptyset))$$

where $w$ is the guidance scale. During training, we randomly drop text conditioning with probability 0.1 to enable this technique.

\textbf{Q-Former Architecture:}

The Q-Former uses learnable query tokens that extract information from text features through attention mechanisms:

\begin{itemize}
    \item \textbf{Learnable Queries}: A small set (2-8) of learned embedding vectors that act as "questions" to ask the language model
    \item \textbf{Self-Attention}: Queries interact with each other to build coherent representations
    \item \textbf{Cross-Attention}: Queries attend to GPT-2 hidden states to extract semantic information relevant for image generation
    \item \textbf{Feed-Forward Networks}: Process the attended features
    \item \textbf{Output Projection}: Project to the conditioning dimension expected by DiT
\end{itemize}

The Q-Former processes text features $\mathbf{F}_{\text{text}} \in \mathbb{R}^{B \times S \times 768}$ through multiple layers:

$$\mathbf{Q}^{(0)} = \text{LearnedEmbedding}(\text{query\_ids})$$
$$\mathbf{Q}^{(l)} = \mathbf{Q}^{(l-1)} + \text{SelfAttn}(\text{LN}(\mathbf{Q}^{(l-1)}))$$
$$\mathbf{Q}^{(l)} = \mathbf{Q}^{(l)} + \text{CrossAttn}(\text{LN}(\mathbf{Q}^{(l)}), \mathbf{F}_{\text{text}})$$
$$\mathbf{Q}^{(l)} = \mathbf{Q}^{(l)} + \text{FFN}(\text{LN}(\mathbf{Q}^{(l)}))$$

where LN denotes LayerNorm. The final queries are projected to produce per-layer conditioning vectors for the DiT model. See Figure \ref{fig:qformer} for details.

\uplevel{\subsection*{Starter Code}}

The files are organized as follows:
\begin{verbatim}
handout/
   requirements.txt
   image_caption_data.py
   train_qformer.py
   eval_qformer.py
   run_in_cloud.ipynb
   dit.py
   ddpm.py
   attention_visualization.py
\end{verbatim}
The codebase is partially built off of \href{https://github.com/cloneofsimo/minRF}{this} open source diffusion implementation. We will use DDPM and the \href{https://www.cs.toronto.edu/~kriz/cifar.html}{CIFAR Dataset} for our implementation.

Below is a description of what you will find in each file. Functions are already implemented and provided for you, unless specified otherwise with TO DO.

\begin{enumerate}
    \item \lstinline{image_caption_data.py}: Functions for the image/caption data and building/caching LLM embeddings
    \begin{itemize}
        \item \lstinline{get_cifar_names}: Gets CIFAR class names 
        \item \lstinline{get_synonyms}: Gets synonyms for CIFAR class names, for data augmentation 
        \item \lstinline{_maybe_apply_synonym}: Possibly swaps out class names for synonyms
        \item \lstinline{get_text_captions}: Retrieves text captions from cached text caption path 
        \item \lstinline{class CIFARTextDataset}: Dataset for training 
        \item \lstinline{prompts_to_padded_hidden_states}: \textbf{TODO:} Function to turn text into GPT-2 embeddings
        \item \lstinline{build_text_states}: Uses \lstinline{get_text_captions} to retrieve and optionally augment text captions with synonyms or dense labels
        \item \lstinline{cache_all_text_states}: Used to cache all GPT-2 embeddings before training for fast and more compute efficient training
    \end{itemize}
    \item \lstinline{train_qformer.py}: Training script
    \begin{itemize}
        \item \lstinline{setup_optimizer_and_scheduler}: \textbf{TODO:} Given a model and a set of arguments, set up the optimizer, scheduler, and loss criterion
        \item \lstinline{maybe_load_resume_states}: Allows for restarting from a given optimizer and weight state in the case of disconnection 
        \item \lstinline{run_qformer_inference}: General inference function
        \item \lstinline{train_cifar10_qformer_dense_captions}: Main training loop
    \end{itemize}
    \item \lstinline{eval_qformer.py}: Evaluation script
    \begin{itemize}
        \item \lstinline{eval_cifar10_qformer}: Main inference code
    \end{itemize}
    \item \lstinline{run_in_cloud.ipynb}: Cloud based ipynb task that has scripts to run training/inference
    \item \lstinline{dit.py}: Model architecture
    \begin{itemize}
        \item \lstinline{DiT_Llama}: Pretrained diffusion transformer 
        \item \lstinline{QueryEmbedder}: \textbf{TODO:} Q-Former implementation (init is provided for you)
        \item \lstinline{Attention}: Multi-head attention with cross-attention support 
        \item \lstinline{FeedForward}: FFN with SwiGLU activation 
    \end{itemize}
    \item \lstinline{ddpm.py}: DDPM pipeline
    \begin{itemize}
        \item \lstinline{DDPM}: forward diffusion and sampling of DDPM
        \item \lstinline{ddpm_schedules}: schedules for DDPM sampling
    \end{itemize}
    \item \lstinline{attention_visualization.py}: Visualization of cross-attention maps
    \item Download links to weights/datasets:\\
    \textbf{TODO:} You should download the three items below. The URLs are in \lstinline{run_in_cloud.ipynb}.
    \begin{itemize}
        \item Pretrained DiT weights: A DiT with our architecture, trained on the class-conditional CIFAR dataset task
        \item Dense CIFAR text captions: A dataset of dense CIFAR captions for enhancing text to image training
        \item Cached GPT-2 text embeddings: Cached GPT-2 embeddings for the dense CIFAR caption dataset
    \end{itemize}
    
\end{enumerate}

\clearpage

\uplevel{\subsection*{Implementation}}

You will implement three key functions that enable Q-Former training and help you understand what it learns:

\begin{enumerate}
    \item \lstinline{prompts_to_padded_hidden_states}: Extract features from a pretrained LLM 
    \item \lstinline{setup_optimizer_and_scheduler}: Configure which parameters train and how
    \item \lstinline{QueryEmbedder.forward}: Implement Q-Former to process text features
    % \item \lstinline{enable_qformer_cross_attention_patch}: Instrument attention layers to save their patterns
\end{enumerate}

All other components (cross attention, training loop, inference) are provided. Your focus is on understanding how efficient fine-tuning works and analyzing the learned attention patterns. 

\vspace{0.3cm}

\noindent\rule{\textwidth}{0.4pt}
\uplevel{\subsection*{Part 1: Extract LLM Features from any layer of GPT-2}}

Implement \lstinline{prompts_to_padded_hidden_states} in \lstinline{shared.py}.

\begin{lstlisting}[language=Python]
    def prompts_to_padded_hidden_states(
    prompts: List[str],
    gpt2,
    tokenizer,
    gpt2_layer_index: int,
    device: torch.device,
):
    """
    Convert a list of text prompts to a single padded tensor 
    of GPT-2 hidden states.

    Returns the hidden states from that specific layer along 
    with masks as a tuple (hidden_states, masks)
    """

\end{lstlisting}

\textit{Your task:}
\begin{itemize}
    \item For all prompts, do the following: 
    \begin{itemize}
        \item Call the tokenizer with \lstinline{return_tensors="pt"} to get valid inputs for GPT-2
        \item Call GPT-2 on these inputs, with \lstinline{output_hidden_states} set to True. You can call \lstinline{outputs.hidden_states} after this to get the hidden states
        \item Given these hidden states, use the \lstinline{gpt2_layer_index} to determine exactly which layer of the hidden states to keep for this prompt
    \end{itemize}
    \item After you have a list of the kept hidden states for all prompts, perform max-length padding of the hidden states, so that the final hidden state size is 
    
    \lstinline{(Batch_size, max_token_length_in_batch, gpt2_hidden_dimension)}.
    
    \item \textbf{NOTE: } Also return a boolean padding tensor of size \lstinline{(Batch_size, max_seq)} where True indicates that a token is not a padded token
\end{itemize}

\uplevel{\subsection*{Part 2: Setup Optimizer and Scheduler}}

Implement \lstinline{setup_optimizer_and_scheduler(model, args)} in \lstinline{train_qformer.py}.

\begin{lstlisting}[language=Python]
def setup_optimizer_and_scheduler(model, args):
    """
    Configure training: freeze/unfreeze parameters, initialize,
    create optimizer/scheduler.
    
    Returns: optimizer, scheduler, criterion
    """
\end{lstlisting}

\textbf{Step 1: Freeze and Unfreeze Parameters} (Parameter-Efficient Training)

The key insight: We want to train only the Q-Former while keeping DiT ($\sim$100M) and GPT-2 ($\sim$117M) frozen. This makes training fast and prevents catastrophic forgetting of pretrained capabilities.

\textit{Your task:}
\begin{itemize}
    \item Set \lstinline{requires_grad = False} for all model parameters, except for those in the query embedder
\end{itemize}

\vspace{0.2cm}

\textbf{Step 2: Create Optimizer}

\textit{Your task:}

Create \lstinline{torch.optim.AdamW} with:
\begin{itemize}
    \item Only parameters where \lstinline{requires_grad=True}
    \item \lstinline{lr=args.lr}, \lstinline{weight_decay=args.weight_decay}, \lstinline{eps=args.adam_epsilon}
    \item \lstinline{betas=(0.9, 0.95)}
\end{itemize}

\vspace{0.2cm}

\textbf{Step 3: Create Learning Rate Scheduler} 

We might expect the first few steps of training to be very unstable, as we're replacing the class embedding from a pretrained model with random noise from a randomly initialized Q-Former. We can add warmup steps to smoothly begin training and determine good momentum and RMS parameters for AdamW.

\textit{Your task:}

If \lstinline{args.warmup_steps > 0}, create a linear warmup scheduler:
\begin{itemize}
    \item Use \lstinline{torch.optim.lr_scheduler.LambdaLR}
    \item The lambda function should return: \lstinline{min(1.0, (step + 1) / warmup_steps)}
    \item This gradually increases learning rate (LR) from 0 to \lstinline{args.lr} over the first \lstinline{warmup_steps}
    \item If \lstinline{warmup_steps <= 0}, set \lstinline{scheduler = None}
\end{itemize}

\textit{Why warmup?} Warmup gradually ramps up the learning rate to prevent unstable updates before model weights and optimizer statistics are well-initialized.

\vspace{0.2cm}

\textbf{Step 4: Create Loss and Return}

\textit{Your task:}

\begin{itemize}
    \item Set criterion to MSE loss
    \item Return designated parameters.
\end{itemize}

\vspace{0.3cm}

\noindent\rule{\textwidth}{0.4pt}


\uplevel{\subsection*{Part 3: Q-Former Forward}}
Implement \lstinline{forward} in \lstinline{./mindiffusion/dit.py}.
\begin{lstlisting}[language=Python]
    def forward(self, cross_attention_states, batch_size=1, 
    cross_attention_mask=None):
    """
        Forward pass of the Q-Former query embedder.

        Returns: Final per-layer features 
        (batch_size, num_dit_layers, conditioning_hidden_size)
    """
\end{lstlisting}

\textit{Your task:}
\begin{itemize}
    \item Get the set of learnable query embeddings from the query table using indices \lstinline{0} to \lstinline{num_queries-1}.
    \item Expand the queries to shape \lstinline{(batch_size, num_queries, transformer_hiden_size)} and put the queries on the same device as the \lstinline{cross_attention_states}.
    \item For each block in \lstinline{self.blocks}, do the following: 
    \begin{itemize}
        \item Apply the self-attention block (i.e. \lstinline{x -> x + self_attn(attn_norm(x))}). 
        \item Apply the cross-attention block \\
(i.e. \lstinline{x -> x + cross_attn(cross_norm(x, cross_attn_states,} \\ \lstinline{cross_attention_mask))}).
        \item Apply the forward-feed block (i.e. \lstinline{x -> x + ffn(ffn_norm(x))}).
        \item You are implementing the equations under Q-Former Architecture.
    \end{itemize}
    \item Use \lstinline{output_proj} to project the resulting shape to the \lstinline{conditioning_hidden_size}.
    \item Map the queries to the per-layer features using the \lstinline{query_to_layer} linear layer to get a tensor with a shape of \lstinline{(batch_size, num_dit_layers, conditioning_hidden_size)}. 
    \item Returns per-layer features.
\end{itemize}
Hint: If you are stuck, look in \lstinline{__init__()} to see what you have to work with. 

\vspace{0.3cm}

\noindent\rule{\textwidth}{0.4pt}

\uplevel{\subsection*{Part 4: Training}}
We train using cached embeddings from GPT-2\footnote{we could have had you generate these, but it would add unnecessary extra time to the already long assignment} with the following text prompts:
\begin{itemize}
    \item simple text prompts with class name: "a photo of a dog"
    \item simple text prompts with a synonym for class name: "a photo of a four legged animal that barks"
    \item dense text prompts with class name: "a photo of a black dog in a grassy field"
    \item dense text prompts with a synonym for class name: "a photo of a four legged animal that barks in a grassy field"
\end{itemize}
These different captions allow the Qformer to learn how to transform semantic text information into dense conditioning information to the pretrained and frozen DiT.
The training loop is provided but you should understand how it works:

\begin{algorithm}[H]
\caption{Q-Former Training (Provided)}
\begin{algorithmic}[1]
\State Load pretrained DiT: \lstinline{model.load_state_dict(torch.load(pretrained_path))}
\State Set up the optimizer, scheduler, and criterion
% \State Freeze DiT: \lstinline{for p in model.parameters(): p.requires_grad = False}
% \State Unfreeze Q-Former: \lstinline{for p in model.query_embedder.parameters(): p.requires_grad = True}
% \State Create optimizer: \lstinline{Adam(model.query_embedder.parameters(), lr=1e-4)}
\For{each epoch}
    \For{each batch $(x_{\text{img}}, y_{\text{caption}})$ from densely labeled, randomly captioned CIFAR-10}
        \State Sample timestep: $t \sim \text{Uniform}(1, T)$
        \State Sample noise: $\epsilon \sim \mathcal{N}(0, I)$
        \State Add noise: $x_t = \sqrt{\bar{\alpha}_t} x_{\text{img}} + \sqrt{1-\bar{\alpha}_t} \epsilon$
        \State Load pre-computed text embedding $y_{\text{emb}}$ for text $y_{\text{caption}})$
        \State \textit{or} calculate $y_{\text{emb}} = \text{extract\_features(gpt2, $y_{\text{caption}}$})$
        \State Predict noise: $\epsilon_{\text{pred}} = \text{DiT}(x_t, t, y_{\text{emb}}, \text{text})$
        \State Compute loss: $\mathcal{L} = \|\epsilon - \epsilon_{\text{pred}}\|^2$
        \State Backpropagate and update only Q-Former parameters
    \EndFor
    \State Generate sample images to monitor training progress
    \State Save Q-Former weights
\EndFor
\end{algorithmic}
\end{algorithm}

\newpage
\textbf{Running Training:} \\
See the \lstinline{run_in_cloud.ipynb} file for the training script. Keep all hyperparameters from these files, but feel free to inspect \lstinline{train_qformer.py} to see all the arguments that you can use and how they work.

\textbf{NOTE: } If your training gets killed after 5 or more epochs, the script will automatically save both your weights and optimizer states every 5 epochs. Use \lstinline{--resume_optimizer_path}, \lstinline{--resume_model_path}, and \lstinline{--start_epoch} to restart from an epoch that is a multiple of 5.

\textbf{Recommended:} Sign in to Google Colab to claim Colab Pro, which is available when you verify with your CMU email. We strongly advise using the A100 GPU when it's available, but if not, using a L4 or T4 should work as well. 


\vspace{0.3cm}

\noindent\rule{\textwidth}{0.4pt}

\textbf{Running Inference:} \\
Using the \lstinline{run_in_cloud.ipynb} file, run \lstinline{eval_qformer.py} in Colab to run inference.

\textbf{NOTE: } The \lstinline{inference.yaml} function has a few examples of how to run inference, but you might have to add more to it to generate all the results for empirical questions.


\textbf{Understanding the Big Picture}

\begin{figure}[h]
\centering
\begin{tabular}{l|l}
\textbf{Your Implementation} & \textbf{What It Enables} \\
\hline
Freeze/unfreeze parameters & Preservation of pretrained model weights and fast training \\
Optimizer \& scheduler & Stable gradient updates \\
QFormer forward & Understanding how the QFormer model works \\
\end{tabular}
\end{figure}

After implementing these functions, the training loop (provided) will:
\begin{enumerate}
    \item Load cached text embeddings from the provided link
    \item Train only the Q-Former using your optimizer
    \item Save samples each epoch for monitoring
\end{enumerate}

\textbf{Key Insight:} You're implementing the ``glue code" that makes efficient multi-modal fine-tuning work. The architecture is provided; your job is to configure \textit{how} it learns and to instrument it for analysis.


\clearpage

\uplevel{\subsection*{Empirical Questions}}

\begin{parts}

\part \textbf{Q-Former Architecture Understanding + Dataset Dive}

\begin{subparts}

\subpart[2] Looking at the code, we observe that each of the Q-Former transformer block has 11.8M parameters, with 2 such blocks. Combining that with the number of parameters of the \textbf{query embeddings} and\textbf{ output projections} how many trainable parameters does the Q-Former have given the config in the training script? (round to nearest million).

Suppose the GPT-2 has 125M parameters, and DiT has 18M parameters, what percentage of the weights in the entire process are trained (not frozen) in our method?

\begin{answer_box}[title=,height=5cm, width=15cm]

\end{answer_box}
\subpart[2] Paste 3 lines from the densely captioned CIFAR dataset. The indices should match with the CIFAR dataset.
\begin{answer_box}[title=,height=5cm, width=15cm]

\end{answer_box}

\end{subparts}

\newpage
\part \textbf{Q-Former Training Understanding}

Consider our Q-Former setup with the following specifications:

\begin{itemize}
    \item the DiT is pretrained on a large corpus of high quality images $I$
    \item the LLM is pretrained on a large corpus of text $T$
    \item the Q-Former training set is on a small set of medium quality images $I'$ with corresponding text $T'$
\end{itemize}

We will consider two scenarios: 

\begin{itemize}
    \item Scenario $A$: Same as in the homework; we only train the Q-Former and keep the LLM and DiT fixed.
    \item Scenario $B$: We will unfreeze parameters of the DiT and LLM along with the Q-Former while training.
\end{itemize}
\begin{subparts}
\subpart[2] Compare and contrast scenario $A$ and $B$ in terms of final asymptotic loss value when trained continuously. Would you expect one scenario to have a lower asymptotic loss, and if so, which one?

\begin{answer_box}[title=,height=5cm, width=15cm]

\end{answer_box}

\subpart[2] Consider running a general text generation validation set (i.e. for a token prediction task) on the LLM used at the end of training in scenarios $A$ and $B$. Which would you expect to have higher mean token accuracy on this general validation set?

\begin{answer_box}[title=,height=5cm, width=15cm]
\end{answer_box}
\end{subparts}

\newpage 
\part \textbf{Q-Former Training}

As indicated in \lstinline{run_in_cloud.ipynb}, train your Q-Former for 50 epochs with \lstinline{num_queries=4}.
[Expected runtime on T4: 2-3 hours] 
[Expected runtime on A100: 1 hour]

\begin{subparts}
\subpart[2] Paste your training Exponential Moving Average (EMA) loss curve for the first 1000 steps. 

\begin{answer_box}[title=,height=6cm, width=15cm]
\end{answer_box}

\subpart[3] Paste the sample grid from epochs 5, 25, and 49 (these should be saved in wandb). How does image quality/prompt adherance improve during training?

\begin{answer_box}[title=,height=10cm, width=15cm]
\end{answer_box}
\end{subparts}

\newpage
\part \textbf{Classifier-Free Guidance}

Experiment with different CFG scales with the original pretrained model: $w \in \{1.0, 2.0, 3.0, 4.0\}$ for the class ``deer" (class 4).

\begin{subparts}
\subpart[2] Paste the 4 images generated with different guidance scales. How does the appearance change as you increase $w$? Which of these actually use the conditional information from the text embeddings?

\begin{answer_box}[title=,height=10cm, width=15cm]
\end{answer_box}

\subpart[1] Based on your results, what seems to be a good guidance scale for this model?

\begin{answer_box}[title=,height=5cm, width=15cm]
\end{answer_box}
\end{subparts}

\newpage
\part \textbf{Text-Conditioned Generation}

After training, generate images from custom text prompts.

\begin{subparts}
\subpart[2] Generate 8 images of cars under two conditions with cfg 3.0:
\begin{itemize}
    \item 8 images with the pretrained class-conditional DiT with class label for automobile
    \item 8 images with the Qformer DiT with text ``a photo of a blue automobile with a bright blue sky"
\end{itemize}
Paste all 8 images. Is there any noticeable difference between the two sets of images?

\begin{answer_box}[title=,height=15cm, width=15cm]
\end{answer_box}

\newpage
\subpart[2] Compare generation of OOD underspecified prompts. Generate 4 photos with caption:``a photo of a vehicle that goes on grass" and compare with 4 photos with caption: ``a photo of a vehicle that goes in water" with cfg 3.0.

Paste the 8 images: How do these images differ? 

\begin{answer_box}[title=,height=10cm, width=15cm]
\end{answer_box}
\end{subparts}

\newpage
\part \textbf{Cross-Attention Maps}

Finally, we will examine the cross-attention maps between the LLM text features and the Q-Former.
The Q-Former uses cross-attention to query the text embeddings. By saving these attention patterns, we can visualize which words the model focuses on when generating specific image features.

\begin{subparts}
\subpart[6] Using the text prompt ``a photo of a red airplane with a green field in the background", generate the attention maps of both layers of the Q-Former at diffusion timesteps $t = 500$ and $t = 100$. Paste these 4 attention maps as well as the final image.

Looking at the attention maps, what do you notice about what tokens are getting attended to and what are getting completely ignored? Why might this be the case?

\textit{Hint:} Take a look at \href{https://arxiv.org/pdf/2309.16588}{Darcet et al. (2023)} and \href{https://arxiv.org/pdf/2309.17453}{Xiao et al. (2023)}. 

\begin{answer_box}[title=,height=9cm, width=15cm]
\end{answer_box}

\end{subparts}
\end{parts}


\clearpage
\sectionquestion{Code Upload}

\begin{parts}

\part[0] Did you upload your code to the appropriate programming slot on Gradescope? \\
\emph{Hint:} The correct answer is `yes'.

    \begin{checkboxes}
     \choice Yes 
     \choice No
    \end{checkboxes}

For this homework, you should upload \lstinline{train_qformer.py}, \lstinline{dit.py}, and \lstinline{image_caption_data.py}.

\end{parts}

\newpage
\sectionquestion{Collaboration Questions}

\begin{parts}

\uplevel{After you have completed all other components of this assignment, report your answers to these questions regarding the collaboration policy. Details of the policy can be found in the syllabus.}

    \part[1] Did you collaborate with anyone on this assignment? If so, list their name or Andrew ID and which problems you worked together on.

        \begin{answer_box}[title=,height=3cm, width=15cm]
        \end{answer_box}

    
    \part[1] Did you find or come across code that implements any part of this assignment? If so, include full details.
        \begin{answer_box}[title=,height=3cm, width=15cm]
        \end{answer_box}
\end{parts}
\end{questions}


\end{document}
