% !Mode:: "TeX:UTF-8"
% !TEX encoding = UTF-8 Unicode

\begin{center}
  \begin{tikzpicture}
    \def\ssep{1cm}
    \def\nodewidth{3*\ssep}
    \def\seg{0.3cm}

    % LLMs as RMs
    \begin{scope}

      \node [anchor=south west,minimum width=3.2*\nodewidth,minimum height=1.3*\ssep,draw,thick,align=center] (llm) at (0,0) {\large{Reward Model (LLM)}};

      \node [anchor=north west,minimum width=0.45*\nodewidth,,minimum height=0.45cm,rounded corners=2pt,draw,fill=gray!20] (input) at ([yshift=-\seg]llm.south west) {\footnotesize{$\mathbf{c}$}};
      \node [anchor=west,minimum width=0.45*\nodewidth,,minimum height=0.45cm,rounded corners=2pt,draw,fill=gray!20] (input2) at ([xshift=6pt]input.east) {\footnotesize{$\mathbf{x}$}};
      \node [anchor=west,minimum width=0.45*\nodewidth,,minimum height=0.45cm,rounded corners=2pt,draw,fill=cyan!20] (input3) at ([xshift=6pt]input2.east) {\footnotesize{$\mathbf{y}_{a}$}};
      \node [anchor=west,minimum width=0.45*\nodewidth,,minimum height=0.45cm,rounded corners=2pt,draw,fill=cyan!20] (input4) at ([xshift=6pt]input3.east) {\footnotesize{$\mathbf{y}_{b}$}};
      \draw [->] ([yshift=1pt]input.north) -- ([yshift=\seg-1pt]input.north);
      \draw [->] ([yshift=1pt]input2.north) -- ([yshift=\seg-1pt]input2.north);
      \draw [->] ([yshift=1pt]input3.north) -- ([yshift=\seg-1pt]input3.north);
      \draw [->] ([yshift=1pt]input4.north) -- ([yshift=\seg-1pt]input4.north);

      \node [anchor=center] (inputlabel) at ([yshift=-0.17cm]input.south) {\scriptsize{prompt}};
      \node [anchor=center] (input2label) at ([yshift=-0.17cm]input2.south) {\scriptsize{input}};
      \node [anchor=center] (input3label) at ([yshift=-0.18cm]input3.south) {\scriptsize{1st output}};
      \node [anchor=center] (input4label) at ([yshift=-0.18cm]input4.south) {\scriptsize{2nd output}};

      \node [anchor=south east,minimum width=0.45cm,minimum height=0.45cm,rounded corners=0pt,draw,fill=orange!20] (output) at ([yshift=\seg,xshift=-0.8*\seg]llm.north east) {\footnotesize{$w$}};
      \draw [->] ([yshift=-\seg+1pt]output.south) -- ([yshift=-1pt]output.south);

      \node [anchor=east] (tokenlabel) at (output.west) {\footnotesize{next token (`A' or `B')}};

      \node [anchor=south west,align=left] (loss) at ([yshift=0.9cm]llm.north west) {\scriptsize{Minimizing the negative probability of token prediction:}\\ \scriptsize{$-\log \mathrm{Pr}_{\phi}(w=\text{A}|[\mathbf{c},\mathbf{x},\mathbf{y}_a,\mathbf{y}_b])$}};
    \end{scope}
  \end{tikzpicture}
\end{center}
