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

\begin{center}
  \begin{tikzpicture}

    \def\ssep{.7cm}
    \def\nsize{0.3cm}

    \tikzstyle{lnode} = [minimum width=7cm,minimum height=1.1cm,inner sep=2pt,draw,thick,fill=white];
    \begin{scope}
      \node [anchor=west] (x0) at (0,0) {\footnotesize{$x_1$}};
      \node [anchor=center] (x1) at ([xshift=\ssep]x0.center) {\footnotesize{$x_2$}};
      \node [anchor=center] (x cdots) at ([xshift=\ssep]x1.center) {\footnotesize{$\cdots$}};
      \node [anchor=center] (x m) at ([xshift=\ssep]x cdots.center) {\footnotesize{$x_m$}};
      \node [anchor=center] (y1) at ([xshift=\ssep]x m.center) {\footnotesize{$y_1$}};
      \node [anchor=center] (y cdots) at ([xshift=\ssep]y1.center) {\footnotesize{$\cdots$}};
      \node [anchor=center] (y n) at ([xshift=\ssep]y cdots.center) {\footnotesize{$y_T$}};
      \node [anchor=north] (y n label) at ([yshift=0.1cm]y n.south) {\scriptsize{(Last Token)}};

      \foreach \x in {x0,x1,x m,y1,y n}
      \draw [->] ([yshift=0.2cm]\x.center) -- ([yshift=0.55cm]\x.center);

      % \node [anchor=center] (ox0) at ([yshift=2.6cm]x0.center) {\footnotesize{$\mathbf{h}_{x_1}$}};
      % \node [anchor=center] (ox1) at ([yshift=2.6cm]x1.center) {\footnotesize{$\mathbf{h}_{x_2}$}};
      % \node [anchor=center] (ox cdots) at ([yshift=2.6cm]x cdots.center) {\footnotesize{$\cdots$}};
      % \node [anchor=center] (ox m) at ([yshift=2.6cm]x m.center) {\footnotesize{$\mathbf{h}_{x_m}$}};
      \node [anchor=center] (oy1) at ([yshift=2.6cm]y1.center) {\footnotesize{$\mathbf{h}_{y_1}$}};
      \node [anchor=center] (oy cdots) at ([yshift=2.6cm]y cdots.center) {\footnotesize{$\cdots$}};
      \node [anchor=center] (oy last) at ([yshift=2.6cm]y n.center) {\footnotesize{$\mathbf{h}_{y_{T}}$}};

      \foreach \x in {oy1,oy last}
      \draw [<-] ([yshift=-0.25cm]\x.center) -- ([yshift=-0.6cm]\x.center);
      
      \node [anchor=south,draw,thick,minimum width=5.2cm,minimum height=1.3cm,fill=white,drop shadow={fill=gray,shadow xshift=.6ex,shadow yshift=-.6ex}] (llm) at ([yshift=.6cm]x m.center) {Reward Model (LLM)};
      \node [anchor=center] (freeze) at ([xshift=-.25cm,yshift=.25cm]llm.south east) {\textcolor{cyan!60}{ \faSnowflake}};
      % \node [anchor=east] (representation) at ([xshift=-0.3cm]ox0.west) {\scriptsize{Representation}};
      % \node [anchor=north west] (representation2) at ([yshift=0.1cm]representation.south west) {\scriptsize{at Each Position}};

      \filldraw [fill=lightgray,draw=white] ([xshift=-0.5cm,yshift=0.6cm]oy last.north west) -- ([xshift=0.5cm,yshift=0.6cm]oy last.north east) -- ([yshift=1.4cm,xshift=0]oy last.north east) -- ([yshift=1.4cm,xshift=0]oy last.north west) -- ([xshift=-0.5cm,yshift=0.6cm]oy last.north west);

      \node [anchor=south] (reward) at ([yshift=1.6cm]oy last.north) {\footnotesize{Reward Model Loss}};
      \node [anchor=south] (Wr) at ([yshift=0.7cm]oy last.north) {\footnotesize{$\mathbf{W}_r$}};
      \node [anchor=center] (not freeze) at ([xshift=1.2em]Wr) {\textcolor{red!60}{\scriptsize \faFire}};
      \node [anchor=east] (linear) at ([xshift=-0.5cm]Wr.west) {\scriptsize{Linear Map}};

      \draw [->] ([yshift=-0.1cm]oy last.north) -- ([yshift=0.58cm]oy last.north);
      \draw [<-] ([yshift=0.1cm]reward.south) -- ([yshift=-0.18cm]reward.south);

      \node (rside) at ([xshift=0.5cm,yshift=-.6cm]oy last.north east|-y n label) {};
      \node (lside) at ([yshift=-.6cm]llm.west|-y n label) {};
      \node at ($(lside)!.5!(rside)$) {\scriptsize{(a)~Parameter Freezing}};
      
    \end{scope}
    
    \begin{scope}[xshift=\textwidth/2]
      \node [anchor=west] (x0) at (0,0) {\footnotesize{$x_1$}};
      \node [anchor=center] (x1) at ([xshift=\ssep]x0.center) {\footnotesize{$x_2$}};
      \node [anchor=center] (x cdots) at ([xshift=\ssep]x1.center) {\footnotesize{$\cdots$}};
      \node [anchor=center] (x m) at ([xshift=\ssep]x cdots.center) {\footnotesize{$x_m$}};
      \node [anchor=center] (y1) at ([xshift=\ssep]x m.center) {\footnotesize{$y_1$}};
      \node [anchor=center] (y cdots) at ([xshift=\ssep]y1.center) {\footnotesize{$\cdots$}};
      \node [anchor=center] (y n) at ([xshift=\ssep]y cdots.center) {\footnotesize{$y_T$}};
      \node [anchor=north] (y n label) at ([yshift=0.1cm]y n.south) {\scriptsize{(Last Token)}};

      \foreach \x in {x0,x1,x m,y1,y n}
      \draw [->] ([yshift=0.2cm]\x.center) -- ([yshift=0.55cm]\x.center);

      \node [anchor=center,minimum width=2.2cm,minimum height=.45cm,draw,densely dashed,rounded corners=2pt] (bg) at ([yshift=2.6cm]y cdots.center) {};

      % \node [anchor=center] (ox0) at ([yshift=2.6cm]x0.center) {\footnotesize{$\mathbf{h}_{x_1}$}};
      % \node [anchor=center] (ox1) at ([yshift=2.6cm]x1.center) {\footnotesize{$\mathbf{h}_{x_2}$}};
      % \node [anchor=center] (ox cdots) at ([yshift=2.6cm]x cdots.center) {\footnotesize{$\cdots$}};
      % \node [anchor=center] (ox m) at ([yshift=2.6cm]x m.center) {\footnotesize{$\mathbf{h}_{x_m}$}};
      \node [anchor=center] (oy1) at ([yshift=2.6cm]y1.center) {\footnotesize{$\mathbf{h}_{y_1}$}};
      \node [anchor=center] (oy cdots) at ([yshift=2.6cm]y cdots.center) {\footnotesize{$\cdots$}};
      \node [anchor=center] (oy last) at ([yshift=2.6cm]y n.center) {\footnotesize{$\mathbf{h}_{y_{T}}$}};

      \foreach \x in {oy1,oy last}
      \draw [<-] ([yshift=-0.25cm]\x.center) -- ([yshift=-0.6cm]\x.center);
      
      \node [anchor=south,draw,minimum width=5.2cm,minimum height=1.3cm,fill=white,drop shadow={fill=gray,shadow xshift=.6ex,shadow yshift=-.6ex},thick] (llm) at ([yshift=.6cm]x m.center) {Reward Model (LLM)};
      % \node [anchor=center] (freeze) at ([xshift=-.25cm,yshift=.25cm]llm.south east) {\textcolor{red!60}{ \faFire}};
      % \node [anchor=east] (representation) at ([xshift=-0.3cm]ox0.west) {\scriptsize{Representation}};
      % \node [anchor=north west] (representation2) at ([yshift=0.1cm]representation.south west) {\scriptsize{at Each Position}};

      \filldraw [fill=lightgray,draw=white] ([xshift=-0.5cm,yshift=0.6cm]oy last.north west) -- ([xshift=0.5cm,yshift=0.6cm]oy last.north east) -- ([yshift=1.4cm,xshift=0]oy last.north east) -- ([yshift=1.4cm,xshift=0]oy last.north west) -- ([xshift=-0.5cm,yshift=0.6cm]oy last.north west);

      \node [anchor=south] (reward) at ([yshift=1.6cm]oy last.north) {\footnotesize{Reward Model Loss}};
      \node [anchor=south] (Wr) at ([yshift=0.7cm]oy last.north) {\footnotesize{$\mathbf{W}_r$}};
      % \node [anchor=center] (not freeze) at ([xshift=1.2em]Wr) {\textcolor{red!60}{\scriptsize \faFire}};
      % \node [anchor=east] (linear) at ([xshift=-0.5cm]Wr.west) {\scriptsize{Linear Map}};

      \node [anchor=south,fill=lightgray,rounded corners=2pt,minimum height=0.8cm,align=center,minimum width=1.9cm,inner sep=0] (ffn softmax) at ([xshift=-2.5cm,yshift=0.6cm]oy last.north west) {\scriptsize{Softmax}\\\scriptsize{Layer}};
      \node [anchor=center,inner sep=0] (sft loss) at (ffn softmax.north|-reward.center) {\footnotesize{SFT Loss}};
      
      % \draw [->,thick,dashed] (oy1.north) .. controls ([yshift=-.2cm]oy1.north|-ffn softmax.east) .. ([yshift=-.2cm]ffn softmax.east);
      \draw [->,densely dashed] (bg.north) .. controls (bg.north|-ffn softmax.south) and (ffn softmax.south|-bg.north) .. (ffn softmax.south);
      \draw [->] ([yshift=.05cm]ffn softmax.north) -- (sft loss.south);
      
      \draw [->] ([yshift=-0.1cm]oy last.north) -- ([yshift=0.58cm]oy last.north);
      \draw [<-] ([yshift=0.1cm]reward.south) -- ([yshift=-0.18cm]reward.south);

      \node (rside) at ([xshift=0.5cm,yshift=-.6cm]oy last.north east|-y n label) {};
      \node (lside) at ([yshift=-.6cm]llm.west|-y n label) {};
      \node at ($(lside)!.5!(rside)$) {\scriptsize{(b)~Regularization}};

    \end{scope}

  \end{tikzpicture}
\end{center}
