LSTMCELL.tex

lstmcell.tex

\documentclass[tikz]{standalone}
\usepackage{amsmath}
\usepackage{bm}
\usepackage{tikz}
\usepackage{pgfplots}
\usetikzlibrary{arrows.meta}
\usetikzlibrary{backgrounds}
\usetikzlibrary{calc}
\usetikzlibrary{fit}
\usetikzlibrary{positioning}
\usetikzlibrary{shapes.arrows}
\usetikzlibrary{shapes.geometric}
\usetikzlibrary{shapes.misc}

\newcommand{\vect}[1]{\bm{#1}}

\tikzset{
    node distance=8mm,
    nn_node_base/.style={
        draw,
        very thick,
        rectangle,
        text=black,
        fill=white,
        inner sep=0pt,
        outer sep=0pt,
        minimum width=11mm,
        minimum height=11mm,
    },
    nn_func_base/.style={
        draw,
        thick,
        circle,
        text=black,
        fill=white,
        inner sep=1pt,
        outer sep=1pt,
        minimum width=8mm,
        minimum height=8mm,
    },
    nn_data_base/.style={
        draw,
        thick,
        text=black,
        fill=white,
        minimum width=10mm,
        minimum height=10mm,
        trapezium,
        trapezium left angle=75,
        trapezium right angle=-75,
        trapezium stretches=true,
        trapezium stretches body=true,
    },
    nn_node_sigmoid/.style={
        nn_node_base,
        color=red!40,
    },
    nn_flow/.style={
        -Latex,
        thin,
        color=black
    },
    nn_dense_vertex/.style={
        circle,
        fill,
        draw,
        minimum size=1.5mm,
    },
}

\pgfplotsset{
    compat=1.18,
    nn_activation_axis/.style={
        samples=71,
        ticks=none,
        xticklabels={\empty},
        yticklabels={\empty},
        axis x line=middle,
        axis y line=middle,
        axis line style={color=black!80,thin,-},
    },
    nn_activation_sigmoid_limits/.style={
        restrict x to domain=-2:2,
        x=10mm/5,
        restrict y to domain=0:1,
        ymin=-.1,ymax=1.1,
        y=10mm/2,
    },
    nn_sigmoid_axis/.style={
        nn_activation_axis,
        nn_activation_sigmoid_limits
    },
    nn_activation_tanh_limits/.style={
        restrict x to domain=-2:2,
        x=10mm/5,
        restrict y to domain=-1:1,
        ymin=-1.1,
        ymax=1.1,
        y=10mm/3,
    },
    nn_tanh_axis/.style={
        nn_activation_axis,
        nn_activation_tanh_limits
    },
    nn_activation_plot/.style={
        black,
        thick
    },
    /pgf/declare function={
        sigmoid(\x) = 1/(1+exp(-2*\x));
    }
}

\newcommand{\nniconsigmoid}{
    \begin{tikzpicture}
        \begin{axis}[nn_sigmoid_axis]
            \addplot[nn_activation_plot] (x,{sigmoid(x)});
        \end{axis}
    \end{tikzpicture}
}

\newcommand{\nnicontanh}{
    \begin{tikzpicture}
        \begin{axis}[nn_tanh_axis]
            \addplot[nn_activation_plot] (x,{tanh(x)});
        \end{axis}
    \end{tikzpicture}
}

\newcommand{\nniconsum}{
   % Like $\oplus$
    \begin{tikzpicture}
        \draw (0,0) node[
            cross out,
            draw=black,
            rotate=45,
            very thick,
            minimum size=3mm,
            inner sep=0pt,
            outer sep=0pt,
        ] {};
    \end{tikzpicture}
}

\newcommand{\nniconmul}{
    \begin{tikzpicture}
        \draw (0,0) node[
            cross out,
            draw=black,
            very thick,
            minimum size=3mm,
        ] {};
    \end{tikzpicture}
}

\newcommand{\nnicondense}{
    \begin{tikzpicture}[node distance=2.7mm]
        \node [nn_dense_vertex] (a0) at (0,0) {};
        \node [nn_dense_vertex] (a1) [below  = of a0] {};
        \node [nn_dense_vertex] (a2) [below  = of a1] {};
        \coordinate (c0) at ($(a0)!0.5!(a1)$);
        \coordinate (c1) at ($(a1)!0.5!(a2)$);
        \node [nn_dense_vertex] (b0) [right = 4mm of c0] {};
        \node [nn_dense_vertex] (b1) [right = 4mm of c1] {};
        \draw [very thin] (a0) edge (b0) edge (b1);
        \draw [very thin] (a1) edge (b0) edge (b1);
        \draw [very thin] (a2) edge (b0) edge (b1);
    \end{tikzpicture}
}

\begin{document}
\begin{tikzpicture}[
    bbox/.style={draw, rounded corners, dashed, inner sep=3mm},
    bboxlabel/.style={rotate=90, anchor=south east},
    xflow/.style={nn_flow,gray},
    hflow/.style={nn_flow,red},
    cflow/.style={nn_flow,green}]
    
    %Data IO
    \coordinate (dataIN_origin) at (-2,-1);
    \node [nn_data_base] (x) at (dataIN_origin) {$\vect{x}$};
    \node [nn_data_base] (htmo) [above = of x] {$\vect{h}^{t-1}$};
    \node [nn_data_base] (ctmo) [above = 45mm of htmo] {$\vect{C}^{t-1}$};
    
    %Forget cell
    \coordinate (forget_origin) at (0,0);
    \node [nn_node_base] (Wf)   [right = of forget_origin] {\nnicondense};
    \node [nn_func_base] (sumF) [above = of Wf] {\nniconsum};
    \node [nn_node_base] (Uf)   [left = of sumF] {\nnicondense};
    \node [nn_node_base] (actF) [above = of sumF] {\nniconsigmoid};
    \draw [nn_flow] (Uf) -- (sumF);
    \draw [nn_flow] (Wf) edge (sumF)
                    (sumF) edge (actF);
    \draw [xflow] (x) -| (Wf);
    \draw [hflow] (htmo) -| (Uf);
    \begin{pgfonlayer}{background}
    \node (forget_box) [bbox, fill=blue!20,fit=(Wf) (Uf) (sumF) (actF), label={[bboxlabel]north west:Forget}] {};
    \end{pgfonlayer}
    
    % Input cell
    \coordinate (input_origin) at (5,0);
    \node [nn_node_base] (Wi)   [right = of input_origin] {\nnicondense};
    \node [nn_func_base] (sumI) [above = of Wi] {\nniconsum};
    \node [nn_node_base] (Ui)   [left = of sumI] {\nnicondense};
    \node [nn_node_base] (actI) [above = of sumI] {\nniconsigmoid};
    \draw [nn_flow] (Ui) -- (sumI);
    \draw [nn_flow] (Wi) edge (sumI)
                    (sumI) edge (actI);
    \draw [xflow] (x) -| (Wi);
    \draw [hflow] (htmo) -| (Ui);
    \begin{pgfonlayer}{background}
    \node (input_box) [bbox, fill=green!20,fit=(Wi) (Ui) (sumI) (actI), label={[bboxlabel]north west:Input}] {};
    \end{pgfonlayer}
    
    % New memory cell
    \coordinate (newmemory_origin) at (10,0);
    \node [nn_node_base] (Wc)   [right = of newmemory_origin] {\nnicondense};
    \node [nn_func_base] (sumC) [above = of Wc] {\nniconsum};
    \node [nn_node_base] (Uc)   [left = of sumC] {\nnicondense};
    \node [nn_node_base] (actC) [above = of sumC] {\nnicontanh};
    \coordinate (midIM)  at ($(actI)!0.5!(actC)$);
    \node [nn_func_base] (mulC) [above =10mm of midIM] {\nniconmul};
    \draw [nn_flow] (Uc) -- (sumC);
    \draw [nn_flow] (Wc) edge (sumC)
                    (sumC) edge (actC);
    \draw [xflow] (x) -| (Wc);
    \draw [hflow] (htmo) -| (Uc);
    \draw [nn_flow] (actI) |- (mulC) node[midway,left] {$\vect{i}^t$};
    \draw [nn_flow] (actC) |- (mulC) node[midway,right] {$\tilde{\vect{C}}^t$};
    \begin{pgfonlayer}{background}
    \node (memory_box) [bbox, fill=yellow!20,fit=(Wc) (Uc) (sumC) (actC), label={[bboxlabel]north west:Memory}] {};
    \end{pgfonlayer}
    
    % Output cell
    \coordinate (output_origin) at (15,0);
    \node [nn_node_base] (Wo)   [right = of output_origin] {\nnicondense};
    \node [nn_func_base] (sumO) [above = of Wo] {\nniconsum};
    \node [nn_node_base] (Uo)   [left = of sumO] {\nnicondense};
    \node [nn_node_base] (actO) [above = of sumO] {\nniconsigmoid};
    \draw [nn_flow] (Uo) -- (sumO);
    \draw [nn_flow] (Wo) edge (sumO)
                    (sumO) edge (actO);
    \draw [xflow] (x) -| (Wo);
    \draw [hflow] (htmo) -| (Uo);
    \begin{pgfonlayer}{background}
    \node (output_box) [bbox, fill=purple!20,fit=(Wo) (Uo) (sumO) (actO), label={[bboxlabel]north west:Output}] {};
    \end{pgfonlayer}
    
    % Short term memory
    \coordinate (short_origin) at (20,0);
    \node [nn_func_base] (mulS) [above = 13mm of short_origin] {\nniconmul};
    \node [nn_node_base] (actS) [above = of mulS] {\nnicontanh};
    \coordinate (midOO) at ($(actO)!0.5!(actS)$);
    \node [nn_data_base] (h) [right = 230mm of htmo] {$\vect{h}^t$};
    \draw [nn_flow] (actO) -- (midOO)
                    (midOO) |- (mulS) node[midway,left] {$\vect{o}^t$};
    \draw [nn_flow] (actS) -- (mulS);
    \draw [hflow] (mulS) |- (h);
    \begin{pgfonlayer}{background}
    \node (short_box) [bbox, fill=red!20,fit=(mulS) (actS) (short_origin), label={[bboxlabel]north west:Short Memory}] {};
    \end{pgfonlayer}
    
    % Long term memory
    \node [nn_func_base] (mulL) at (ctmo -| actF) {\nniconmul};
    \node [nn_func_base] (sumL) at (mulL -| mulC) {\nniconsum};
    \node [nn_data_base] (c) at (sumL -| actS) {$\vect{C}^t$};
    \node (CC) [right = 20mm of c] {};
    \draw [cflow] (c) edge (actS)
                  (c) -- (CC);
    \draw [nn_flow] (actF) -- (mulL) node[midway,left] {$\vect{f}^t$};
    \draw [nn_flow] (mulC) edge (sumL);
    \draw [cflow] (ctmo) edge (mulL)
                  (mulL) edge (sumL)
                  (sumL) edge (c);
                  
\end{tikzpicture}

\end{document}