\section{Introduction}

During the last decade, deep learning \citep{goodfellow2016dlbook}, the study of neural networks (NNs), has achieved ground-breaking results in diverse areas such as computer vision \citep{krizhevsky12alexnet,he2016resnet,long15fully,chen2018deeplab}, natural language processing \citep{hinton2012speech,vaswani2017attention,devlin2019bert}, generative modeling \citep{kingma2014vae,goodfellow2014gan} and reinforcement learning \citep{mnih2015dqn,silver2016alphago}, and various network designs have been proposed. However, neural networks have been treated largely as ``black-box'' function approximators, and their designs have chiefly been found via trial-and-error, with little or no theoretical justification. A major cause that hinders the theoretical analysis is the current overly generic modeling of neural networks as function approximators: simply interpreting a neural network as a composition of parametrized functions provides little insight to decipher the nature of its components or its behavior during the learning process.

In this paper, we show that a neural network can actually be interpreted as a probabilistic graphical model (PGM) called chain graph (CG) \citep{koller2009pgmbook}, and feed-forward as an efficient approximate probabilistic inference on it. This offers specific interpretations for various neural network components, allowing for in-depth theoretical analysis and derivation of new approaches.

\begin{figure}[t]
\centering
\includegraphics[width=0.39\linewidth]{banner}
\hspace{0.02\linewidth}
\includegraphics[width=0.39\linewidth]{activation}
\caption{Neural networks can be interpreted as layered chain graphs where activation functions are determined by node distributions. \textsl{Left}: An example neural network interpreted as a chain graph with three chain components which represent its layers; \textsl{Right}: A variety of activation functions (softplus, ReLU, leaky ReLU) approximated by nodes following rectified Gaussian distributions ($e,q$ as in Eq.~\eqref{eq:ff}). We visualize the approximations stochastically by averaging over 200 samples.} \label{fig:banner}
\end{figure}

\subsection{Related work} \label{subsec:relatedwork}

In terms of theoretical understanding of neural networks, a well known result based on the function approximator view is the universal approximation theorem \citep{goodfellow2016dlbook}, however it only establishes the representational power of NNs. Also, there have been many efforts on alternative NN interpretations. One prominent approach identifies infinite width NNs as Gaussian processes \citep{neal1996phd,lee2018dnnasgp}, enabling kernel method analysis \citep{jacot2018ntk}. Other works also employ theories such as optimal transport \citep{genevay2017gan,chizat2018global} or mean field \citep{mei2019mean}. These approaches lead to interesting findings, however they tend to only hold under limited or unrealistic settings and have difficulties interpreting practical real-world NNs.

Alternatively, some existing works study the post-hoc interpretability \citep{lipton2018mythos}, proposing methods to analyze the empirical behavior of trained neural networks: activation maximization \citep{erhan2009visualizing}, typical input synthesis \citep{nguyen2016synthinput}, deconvolution \citep{zeiler2014deconv},  layer-wise relevance propagation \citep{bach2015lrp}, etc. These methods can offer valuable insights to the practical behavior of  neural networks, however they represent distinct approaches and focuses, and are all limited within the function approximator view.

Our work links neural networks to probabilistic graphical models \citep{koller2009pgmbook}, a rich theoretical framework that models and visualizes probabilistic systems composed of random variables (RVs) and their interdependencies. There are several types of graphical models. The chain graph  model \citep{koller2009pgmbook} used in our work is a general form that unites directed and undirected variants, visualized as a partially directed acyclic graph (PDAG). Interestingly, there exists a series of works on constructing hierarchical graphical models for data-driven learning problems, such as sigmoid belief network \citep{neal1992sbn}, deep belief network \citep{hinton2006dbn}, deep Boltzmann machine \citep{salakhutdinov2012dbm} and sum product network \citep{poon2011spn}. As alternatives to neural networks, these models have shown promising potentials for generative modeling and unsupervised learning. Nevertheless, they are yet to demonstrate competitive performances over neural network for discriminative learning.

Neural networks and graphical models have so far been treated as two distinct approaches in general. Existing works that combine them \citep{zheng2015crfasrnn, chen2018deeplab, lampe2016crfnlp} treat either neural networks as function approximators for amortized inference, or graphical models as post-processing steps. Some consider neural networks as graphical models with deterministic hidden nodes \citep{buntine1994gmasnn}. However this is an atypical degenerate regime. To the best of our knowledge, our work provides the first rigorous and comprehensive formulation of a (non-degenerate) graphical model interpretation for neural networks in practical use.

\subsection{Our contributions}

The main contributions of our work are summarized as follows:
\begin{itemize}
\item We propose a layered chain graph representation of neural networks, interpret feed-forward as an approximate probabilistic inference procedure, and show that this interpretation provides an extensive coverage of practical NN components (Section~\ref{sec:nnascg});
\item To illustrate its advantages, we show with concrete examples (residual block, RNN, dropout) that the chain graph interpretation enables coherent and in-depth theoretical support,  and provides additional insights to various empirically established network structures (Section~\ref{sec:case_studies});
\item Furthermore, we demonstrate the potential of the chain graph interpretation for discovering new approaches by using it to derive a novel stochastic inference method named partially collapsed feed-forward, and establish experimentally its empirical effectiveness (Section~\ref{sec:pcff}).
\end{itemize}

\section{Chain graph interpretation of neural networks}
\label{sec:nnascg}

Without further delay, we derive the chain graph interpretation of neural networks in this section. We will state and discuss the main results here and leave the proofs in the appendix.

\subsection{The layered chain graph representation}
\label{subsec:cgrepr}

We start by formulating the so called \textsl{layered chain graph} that corresponds to neural networks we use in practice: Consider a system represented by $L$ layers of random variables $(\mathbf{X}^{1}, \dots, \mathbf{X}^{L})$, where $X^{l}_i$ is the $i$-th variable node in the $l$-th layer, and denote $N^l$ the number of nodes in layer $l$. We assume that nodes $X^{l}_i$ in the same layer $l$ have the same distribution type characterized by a feature function $\mathbf{T}^l$ that can be multidimensional. Also, we assume that the layers are ordered topologically and denote $Pa(\mathbf{X}^l)$ the parent layers of $\mathbf{X}^l$. To ease our discussion, we assume that $\mathbf{X}^1$ is the input layer and $\mathbf{X}^L$ the output layer (our formulation can easily extend to multi-input/output cases). A layered chain graph is then defined as follows:

\begin{atdef}
A \textsl{layered chain graph} that involves $L$ layers of random variables $(\mathbf{X}^{1}, \dots, \mathbf{X}^{L})$ is a chain graph that encodes the overall distribution $P(\mathbf{X}^2, \dots, \mathbf{X}^L | \mathbf{X}^1)$ such that:
\begin{enumerate}
\item It can be factored into layerwise chain components $P(\mathbf{X}^{l}|Pa(\mathbf{X}^l))$ following the topological order, and nodes $X^{l}_i$ within each chain component $P(\mathbf{X}^{l}|Pa(\mathbf{X}^l))$  are conditionally independent given their parents (this results in bipartite chain components), thus allowing for further decomposition into nodewise conditional distributions $P(X^{l}_i|Pa(\mathbf{X}^l))$ . This means we have
\begin{equation}
P(\mathbf{X}^2, \dots, \mathbf{X}^L | \mathbf{X}^1) = \prod_{l=2}^{L} P(\mathbf{X}^{l}|Pa(\mathbf{X}^l)) = \prod_{l=2}^{L} \prod_{i=1}^{N^{l}} P(X_i^{l}|Pa(\mathbf{X}^l));
\end{equation}

\item For each layer $l$ with parent layers $Pa(\mathbf{X}^l)=\{ \mathbf{X}^{p_1}, \dots \mathbf{X}^{p_n} \}, p_1,\dots,p_n \in \{ 1,\dots,l-1 \}$, its nodewise conditional distributions $P(X^{l}_i|Pa(\mathbf{X}^l))$ are modeled by pairwise conditional random fields (CRFs) with with unary ($\mathbf{b}^{l}_i$) and pairwise ($\mathbf{W}_{j, i}^{p, l}$) weights (as we will see, they actually correspond to biases and weights in NN layers):
\begin{align}
&P(X^{l}_i|Pa(\mathbf{X}^l)) = f^{l} \big( \mathbf{T}^l(X^{l}_i), \mathbf{e}^{l}_i\big(\mathbf{T}^{p_1}(\mathbf{X}^{p_1}),\dots,\mathbf{T}^{p_n}(\mathbf{X}^{p_n})\big) \big) \label{eq:activgeneral} \\
\text{with}\quad &\mathbf{e}^{l}_i\big(\mathbf{T}^{p_1}(\mathbf{X}^{p_1}),\dots,\mathbf{T}^{p_n}(\mathbf{X}^{p_n})\big) = \mathbf{b}^{l}_i + \sum_{p=p_1}^{p_n} \sum_{j=1}^{N^{p}} \mathbf{W}_{j, i}^{p, l} \mathbf{T}^{p}(X^{p}_j). \label{eq:activgeneralpreact}
\end{align}
\end{enumerate}
\end{atdef}

Figure~\ref{fig:banner}~Left illustrates an example three-layer network as layered chain graph and its chain component factorization. For exponential family distributions \citep{koller2009pgmbook}, the general form in Eq.~\eqref{eq:activgeneral} simply becomes $P(X^{l}_i|Pa(\mathbf{X}^l)) \propto \exp \big( \mathbf{T}^l(X^{l}_i) \cdot \mathbf{e}^{l}_i\big(\mathbf{T}^{p_1}(\mathbf{X}^{p_1}),\dots,\mathbf{T}^{p_n}(\mathbf{X}^{p_n})\big) \big)$.

\subsection{Feed-forward as approximate probabilistic inference} \label{subsec:ffprobainfer}

To identify layered chain graphs with real-world neural networks, we need to show that they can behave the same way during inference and learning. For this, we establish the fact that feed-forward can actually be seen as performing an approximate probabilistic inference on a layered chain graph: Given an input sample $\tilde{\mathbf{x}}^1$, we consider the problem of inferring the marginal distribution $Q^l_i$ of a node $X^l_i$ and its expected features $\mathbf{q}^l_i$, defined as
\begin{equation}
Q^l_i (x^l_i | \tilde{\mathbf{x}}^1) = P(X^l_i=x^l_i | \mathbf{X}^1=\tilde{\mathbf{x}}^1);\quad \mathbf{q}^l_i = \mathbb{E}_{Q^l_i}[\mathbf{T}^l(X^l_i)] \enspace (
\mathbf{q}^1 = \tilde{\mathbf{x}}^1). \label{eq:qdefs}
\end{equation}
Consider a non-input layer $l$ with parent layers $p_1,\dots,p_n$, the independence assumptions encoded by the layered chain graph lead to the following recursive expression for marginal distributions $Q$:
\begin{equation}
Q^l_i (x^l_i | \tilde{\mathbf{x}}^1) =  \mathbb{E}_{Q^{p_1},\dots,Q^{p_n}}[P(x^l_i | Pa(\mathbf{X}^l))]. \label{eq:qrecursive}
\end{equation}
However, the above expression is in general intractable, as it integrates over the entire admissible states of all parents nodes in $Pa(\mathbf{X}^l)$. To proceed further, simplifying approximations are needed. Interestingly, by using linear approximations, we can obtain the following results (in case of discrete random variable the integration in Eq.~\ref{eq:ff} is replaced by summation):

\begin{atprop} \label{thm:ff}
If we make the assumptions that the corresponding expressions are approximately linear w.r.t.\ parent features $\mathbf{T}^{p_1}(\mathbf{X}^{p_1}),\dots,\mathbf{T}^{p_n}(\mathbf{X}^{p_n})$, we obtain the following approximations:
\begin{align}
&Q^l_i (x^l_i | \tilde{\mathbf{x}}^1) \approx f^{l} \big( \mathbf{T}^l(x^{l}_i), \mathbf{e}^{l}_i(\mathbf{q}^{p_1},\dots,\mathbf{q}^{p_n}) \big); \label{eq:distff} \\
&\mathbf{q}^l_i \approx \int_{x^l_i} \mathbf{T}^l(x^l_i) f^{l} \big( \mathbf{T}^l(x^{l}_i), \mathbf{e}^{l}_i(\mathbf{q}^{p_1},\dots,\mathbf{q}^{p_n}) \big) dx^l_i := \mathbf{g}^l(\mathbf{e}^{l}_i(\mathbf{q}^{p_1},\dots,\mathbf{q}^{p_n})). \label{eq:ff}
\end{align}
Especially, Eq.~\eqref{eq:ff} is a feed-forward expression for expected features $\mathbf{q}^l_i$ with activation function $\mathbf{g}^l$ determined by $\mathbf{T}^l$ and $f^l$, i.e.\ the distribution type of random variable nodes in layer $l$.
\end{atprop}

The proof is provided in Appendix~\ref{asubsec:proof_thm_ff}. This allows us to identify feed-forward as an approximate probabilistic inference procedure for layered chain graphs. For learning, the loss function is typically a function of $(Q^L,\mathbf{q}^L)$ obtainable via feed-forward, and we can follow the same classical neural network parameter update using stochastic gradient descent and backpropagation. Thus we are able to replicate the exact neural network training process with this layered chain graph framework.

The following corollary provides concrete examples of some common activation functions $\mathbf{g}$ (we emphasize their names in bold, detailed formulations and proofs are given in Appendix~\ref{asubsec:proof_cor_activationfuncs}):

\begin{atcor} \label{cor:activationfuncs}
We have the following node distribution \-- activation function correspondences:
\begin{enumerate}
\item Binary nodes taking values $\{\alpha, \beta\}$ results in sigmoidal activations, especially, we obtain \textbf{sigmoid} with $\alpha=0, \beta=1$ and \textbf{tanh} with $\alpha=-1, \beta=1$ ($\alpha, \beta$ are interchangeable);

\item Multilabel nodes characterized by label indicator features result in the \textbf{softmax} activation;

\item Variants of (leaky) rectified Gaussian distributions ($T^{l}_i(X^{l}_i) = X^{l}_i = \max(\epsilon Y^{l}_i, Y^{l}_i)$ with $Y^{l}_i \sim \mathcal{N} \big( e^{l}_i, (s^l_i(e^{l}_i))^2 \big)$) can approximate activations such as \textbf{softplus} ($\epsilon=0, s^l_i \approx 1.7761$) and \textbf{$\epsilon$-leaky rectified linear unit} (ReLU) ($s^l_i =\tanh (e^{l}_i)$) including \textbf{ReLU} ($\epsilon=0$) and \textbf{identity} ($\epsilon=1$).
\end{enumerate}
\end{atcor}

Figure~\ref{fig:banner}~Right illustrates activation functions approximated by various rectified Gaussian variants. We also plotted (in orange) an alternative approximation of ReLU with sigmoid-modulated standard deviation proposed by \citet{nair2010relu} which is less accurate around the kink at the origin.

The linear approximations, needed for feed-forward, is coarse and only accurate for small pairwise weights ($\|\mathbf{W}\| \ll 1$) or already linear regions. This might justify weight decay beyond the general ``anti-overfit'' argument and the empirical superiority of piecewise linear activations like ReLU \citep{nair2010relu}.  Conversely, as a source of error, it might explain some ``failure cases'' of neural networks such as their vulnerability against adversarial samples, see e.g.,\ \citet{goodfellow2015adversarial}.

\subsection{Generality of the chain graph interpretation} \label{subsec:gen_nnascg}

The chain graph interpretation formulated in Sections~\ref{subsec:cgrepr} and \ref{subsec:ffprobainfer} is a general framework that can describe many practical network structures. To demonstrate this, we list here a wide range of neural network designs (marked in bold) that are chain graph interpretable.

\begin{itemize}
\item In terms of network architecture, it is clear that the chain graph interpretation can model networks of arbitrary depth, and with general multi-branched structures  such as \textbf{inception modules} \citep{szegedy2015googlenet} or \textbf{residual blocks} \citep{he2016resnet, he2016resnet2} discussed in Section~\ref{subsec:residual}. Also, it is possible to built up \textbf{recurrent neural networks (RNNs)} for sequential data learning, as we will see in Section~\ref{subsec:rnn}. Furthermore,  the modularity of chain components justifies \textbf{transfer learning via partial reuse of pre-trained networks}, e.g.,\ backbones trained for image classification can be reused for segmentation \citep{chen2018deeplab}.

\item In terms of layer structure, we are free to employ sparse connection patterns and shared/fixed weight, so that we can obtain not only \textbf{dense connections}, but also connections like \textbf{convolution}, \textbf{average pooling} or \textbf{skip connections}. Moreover, as shown in Section~\ref{subsec:dropout}, \textbf{dropout} can be reproduced by introducing and sampling from auxiliary random variables, and normalization layers like \textbf{batch normalization} \citep{ioffe2015batchnorm} can be seen as reparametrizations of node distributions and fall within the general form (Eq.~\eqref{eq:activgeneral}). Finally, we can extend the layered chain graph model to allow for intra-layer connections, which enables \textbf{non-bipartite CRF layers} which are typically used on output layers for structured prediction tasks like image segmentation \citep{zheng2015crfasrnn, chen2018deeplab} or named entity recognition \citep{lampe2016crfnlp}. However, feed-forward is no longer applicable through these intra-connected layers.

\item Node distributions can be chosen freely, leading to a variety of nonlinearities (e.g.,\ Corollary~\ref{cor:activationfuncs}).
\end{itemize}

\section{Selected case studies of existing neural network designs} \label{sec:case_studies}

The proposed chain graph interpretation offers a detailed description of the underlying mechanism of neural networks. This allows us to obtain novel theoretical support and insights for various network designs which are consistent within a unified framework. We illustrate this with the following concrete examples where we perform in-depth analysis based on the chain graph formulation.

\subsection{Residual block as refinement module} \label{subsec:residual}

The residual block, proposed originally in \citet{he2016resnet} and improved later \citep{he2016resnet2} with the preactivation form, is an effective design for building up very deep networks. Here we show that a preactivation residual block corresponds to a refinement module within a chain graph. We use \textsl{modules} to refer to encapsulations of layered chain subgraphs as input--output mappings without specifying their internal structures. A refinement module is defined as follows:

\begin{atdef} \label{def:refinementmodule}
Given a base submodule from layer $\mathbf{X}^{l-1}$ to layer $\mathbf{X}^l$, a \textsl{refinement module} augments this base submodule with a side branch that chains a copy of the base submodule (sharing weight with its original) from $\mathbf{X}^{l-1}$ to a duplicated layer $\tilde{\mathbf{X}}^l$, and then a refining submodule from $\tilde{\mathbf{X}}^l$ to $\mathbf{X}^l$.
\end{atdef}

\begin{figure}[ht]
\centering
\includegraphics[width=0.24\linewidth]{resnet_chaingraph}
\hspace{0.02\linewidth}
\includegraphics[width=0.6\linewidth]{resnet_compgraph2}
\caption{Example of a refinement module (left) and its corresponding computational graph (right), composed of a base submodule $X^{l-1} \to X^l$ (blue background) and a refining submodule $\tilde{X}^l \to Z^l \to X^l$ (red background). In the computational graph each $W,b$ represents a linear connection (Eq.~\eqref{eq:activgeneralpreact}) and $\sigma$ an activation function. Same color identifies corresponding parts in the two graphs. We see that this refinement module corresponds exactly to a preactivation residual block.} \label{fig:residualblock}
\end{figure}

\begin{atprop} \label{prop:residual}
A refinement module corresponds to a preactivation residual block.
\end{atprop}

We provide a proof in Appendix~\ref{asubsec:proof_prop_residual} and illustrate this correspondence in Figure~\ref{fig:residualblock}. An interesting remark is that the refinement process can be recursive: the base submodule of a refinement module can be a refinement module itself. This results in a sequence of consecutive residual blocks.

While a vanilla layered chain component encodes a generalized linear model during feed-forward (c.f.\ Eqs.~\eqref{eq:ff},\eqref{eq:activgeneralpreact}), the refinement process introduces a nonlinear extension term to the previously linear output preactivation, effectively increasing the representational power. This provides a possible explanation to the empirical improvement generally observed when using residual blocks.

Note that it is also possible to interpret the original postactivation residual blocks, however in a somewhat artificial manner, as it requires defining identity connections with manually fixed weights.

\subsection{Recurrent neural networks} \label{subsec:rnn}

Recurrent neural networks (RNNs) \citep{goodfellow2016dlbook} are widely used for handling sequential data. An unrolled recurrent neural network can be interpreted as a dynamic layered chain graph constructed as follows: a given base layered chain graph is copied for each time step, then these copies are connected together through recurrent chain components following the Markov assumption \citep{koller2009pgmbook}: each recurrent layer $\mathbf{X}^{l,t}$ at time $t$ is connected by its corresponding layer $\mathbf{X}^{l,t-1}$ from the previous time step $t-1$. Especially, denoting $Pa^t(\mathbf{X}^{l,t})$ the non-recurrent parent layers of $\mathbf{X}^{l,t}$ in the base chain graph, we can easily interpret the following two variants:

\begin{atprop} \label{prop:rnn} Given a recurrent chain component that encodes $P(\mathbf{X}^{l,t}|Pa^t(\mathbf{X}^{l,t}), \mathbf{X}^{l,t-1})$,
\begin{enumerate}
\item It corresponds to a simple (or vanilla / Elman) recurrent layer \citep{goodfellow2016dlbook} if the connection from $\mathbf{X}^{l,t-1}$ to $\mathbf{X}^{l,t}$ is dense;
\item It corresponds to an independently RNN (IndRNN) \citep{li2018indrnn} layer if the conditional independence assumptions among the nodes $X^{l,t}_i$ within layer $l$ are kept through time:
\begin{equation}
\forall i \in \{1,\dots, N^l\}, \ P(X^{l,t}_i|Pa^t(\mathbf{X}^{l,t}), \mathbf{X}^{l,t-1}) = P(X^{l,t}_i|Pa^t(\mathbf{X}^{l,t}), X^{l,t-1}_i).
\end{equation}
\end{enumerate}
\end{atprop}

We provide a proof in Appendix~\ref{asubsec:proof_prop_rnn} and illustrates both variants in Figure~\ref{fig:rnn}.

\begin{figure}[ht]
\centering
\includegraphics[width=0.35\linewidth]{srn}
\hspace{0.06\linewidth}
\includegraphics[width=0.35\linewidth]{indrnn}
\caption{Comparison of an simple recurrent layer (left) v.s.\ an IndRNN (right) recurrent layer. IndRNN, the better variant, enforces the intra-layer conditional independence through time.} \label{fig:rnn}
\end{figure}

The simple recurrent layer, despite its exhaustive dense recurrent connection, is known to suffer from vanishing/exploding gradient and can not handle long sequences. The commonly used long-short term memory \citep{hochreiter1997lstm} and gated recurrent unit \citep{cho2014gru} alleviate this issue via long term memory cells and gating. However, they tend to result in bloated structures, and still cannot handle very long sequences \citep{li2018indrnn}. On the other hand, IndRNNs can process much longer sequences and significantly outperform not only simple RNNs, but also LSTM-based variants \citep{li2018indrnn,li2019indrnn}. This indicates that the assumption of intra-layer conditional independence through time, analogue to the local receptive fields of convolutional neural networks, could be an essential sparse network design tailored for sequential modeling.

\subsection{Dropout} \label{subsec:dropout}

Dropout \citep{srivastava2014dropout} is a practical stochastic regularization method commonly used especially for regularizing fully connected layers. As we see in the following proposition, from the chain graph point of view, dropout corresponds to introducing Bernoulli auxiliary random variables that serve as noise generators for feed-forward during training:

\begin{atprop} \label{prop:dropout}
Adding dropout with drop rate $1 - p^l$ to layer $l$ corresponds to the following chain graph construction: for each node $X^l_i$ in layer $l$ we introduce an auxiliary Bernoulli random variable $D^l_i \sim \text{Bernoulli}(p^l)$ and multiply it with the pairwise interaction terms in all preactivations (Eq.~\eqref{eq:activgeneralpreact}) involving $X^l_i$ as parent (this makes $D^l_i$ a parent of all child nodes of $X^l_i$ and extend their pairwise interactions with $X^l_i$ to ternary ones). The behavior of dropout is reproduced exactly if:
\begin{itemize}
\item During training, we sample auxiliary nodes $D^l_i$ during each feed-forward. This results in dropping each activation $\mathbf{q}^l_i$ of node $X^l_i$ with probability $1 - p^l$;
\item At test time, we marginalize auxiliary nodes $D^l_i$ during each feed-forward. This leads to deterministic evaluations with a constant scaling of $p^l$ for the node activations $\mathbf{q}^l_i$.
\end{itemize}
\end{atprop}

We provide a proof in Appendix~\ref{asubsec:proof_prop_dropout}. Note that among other things, this chain graph interpretation of dropout provides a theoretical justification of the constant scaling at test time. This was originally proposed as a heuristic in \citet{srivastava2014dropout} to maintain consistent behavior after training.

\section{Partially collapsed feed-forward} \label{sec:pcff}

The theoretical formulation provided by the chain graph interpretation can also be used to derive new approaches for neural networks. It allows us to create new deep learning methods following a coherent framework that provides specific semantics to the building blocks of neural networks. Moreover, we can make use of the abundant existing work from the PGM field, which also serves as a rich source of inspiration. As a concrete example, we derive in this section a new stochastic inference procedure called partially collapsed feed-forward (PCFF) using the chain graph formulation.

\subsection{PCFF: chain graph formulation} \label{subsec:pcffcg}

A layered chain graph, which can represent a neural network, is itself a probabilistic graphical model that encodes an overall distribution conditioned on the input. This means that, to achieve stochastic behavior, we can directly draw samples from this distribution, instead of introducing additional ``noise generators'' like in dropout. In fact, given the globally directed structure of layered chain graph, and the fact that the conditioned input nodes are ancestral nodes without parent, it is a well-known PGM result that we can apply forward sampling (or ancestral sampling) \citep{koller2009pgmbook} to efficiently generate samples: given an input sample $\tilde{\mathbf{x}}^1$, we follow the topological order and sample each non-input node $X^l_i$ using its nodewise distribution (Eq.~\eqref{eq:activgeneral}) conditioned on the samples $(\mathbf{x}^{p_1}, \dots,\mathbf{x}^{p_n})$ of its parents. Compared to feed-forward, forward sampling also performs a single forward pass, but generates instead an unbiased stochastic sample estimate.

While in general an unbiased estimate is preferable and the stochastic behavior can also introduce regularization during training \citep{srivastava2014dropout}, forward sampling can not directly replace feed-forward, since the sampling operation is not differentiable and will jeopardize the gradient flow during backpropagation. To tackle this, one idea is to apply the reparametrization trick \citep{kingma2014vae} on continuous random variables (for discrete RVs the Gumbel softmax trick \citep{jang2017gumbelsoftmax} can be used but requires additional continuous relaxation). An alternative solution is to only sample part of the nodes as in the case of dropout.

The proposed partially collapse feed-forward follows the second idea: we simply ``mix up'' feed-forward and forward sampling, so that for each forward inference during training, we randomly select a portion of nodes to sample and the rest to compute deterministically with feed-forward. Thus for a node $X_i^l$ with parents $(\mathbf{X}^{p_1},\dots,\mathbf{X}^{p_n})$, its forward inference update becomes
\begin{equation}
\mathbf{q}_i^l \gets
\begin{cases}
\mathbf{g}^l(\mathbf{e}^l_i(\mathbf{q}^{p_1}, \dots, \mathbf{q}^{p_n})) &\text{ if collapsed (feed-forward)}; \\
\mathbf{T}^l(x_i^l),\ x_i^l \sim f^{l} \big( \mathbf{T}^l(X^{l}_i), \mathbf{e}^{l}_i(\mathbf{q}^{p_1},\dots,\mathbf{q}^{p_n}) \big) &\text{ if uncollapsed (forward sampling)}.
\end{cases}
\end{equation}
Following the collapsed sampling \citep{koller2009pgmbook} terminology, we call this method the partially collapsed feed-forward (PCFF). PCFF is a generalization over feed-forward and forward sampling, which can be seen as its fully collapsed / uncollapsed extremes. Furthermore, it offers a bias--variance trade-off, and can be combined with the reparametrization trick to achieve unbiased estimates with full sampling, while simultaneously maintaining the gradient flow.

\subsection{PCFF: experimental validation} \label{subsec:pcffexp}

In the previous sections, we have been discussing existing approaches whose empirical evaluations have been thoroughly covered by prior work. The novel PCFF approach proposed in this section, however, requires experiments to check its practical effectiveness. For this we conduct here a series of experiments\footnote{Implementation available at: \url{https://github.com/tum-vision/nnascg}}. Our emphasis is to understand the behavior of PCFF under various contexts and not to achieve best result for any specific task. We only use chain graph interpretable components, and we adopt the reparameterization trick  \citep{kingma2014vae} for ReLU PCFF samples.

The following experiments show that PCFF is overall an effective stochastic regularization method. Compared to dropout, it tends to produce more consistent performance improvement, and can sometimes outperform dropout. This confirms that our chain graph based reasoning has successfully found an interesting novel deep learning method.

\begin{figure}[ht]
\centering
\includegraphics[width=0.47\linewidth]{mnist_}
\hspace{0.04\linewidth}
\includegraphics[width=0.36\linewidth]{cifar}
\caption{Comparison of stochastic methods (None/Dropout/PCFF) in terms of image classification test errors (lower is better) under various settings. Left: MNIST/FashionMNIST datasets with a simple dense network and tanh/ReLU activation functions; Right: CIFAR-10 dataset with ResNet20 and varying drop/sample rates. All reported results are average values of three runs. Compared to dropout, PCFF can achieve comparable results, and tend to deliver more consistent improvements. } \label{fig:pcffexp}
\end{figure}

\paragraph{Simple dense network}
We start with a simple network with two dense hidden layers of 1024 nodes to classify MNIST \citep{lucun1998mnist} and FashionMNIST \citep{xiao2017fashionmnist} images. We use PyTorch \citep{paszke2017pytorch}, train with stochastic gradient descent (learning rate $0.01$, momentum $0.9$), and set up 20\% of training data as validation set for performance monitoring and early-stopping. We set drop rate to 0.5 for dropout, and for PCFF we set the sample rate to 0.4 for tanh and 1.0 (full sampling) for ReLU. Figure~\ref{fig:pcffexp}~Left reports the test errors with different activation functions and stochastic regularizations.

We see that dropout and PCFF are overall comparable, and both improve the results in most cases. Also, the ReLU activation consistently produces better results that tanh. Additional experiments show that PCFF and dropout can be used together, which sometimes yields improved performance.

\paragraph{Convolutional residual network}
To figure out the applicability of PCFF in convolutional residual networks, we experiment on CIFAR-10 \citep{krizhevsky09cifar} image classification. For this we adapt an existing implementation
\citep{Idelbayev18a} to use the preactivation variant. We focus on the ResNet20 structure, and follow the original learning rate schedule except for setting up a validation set of 10\% training data to monitor training performance. Figure~\ref{fig:pcffexp}~Right summarizes the test errors under different drop/sample rates.

We observe that in this case PCFF can improve the performance over a wide range of sample rates, whereas dropout is only effective with drop rate $0.1$, and large drop rates in this case significantly deteriorate the performance. We also observe a clear trade-off of the PCFF sample rate, where a partial sampling of 0.3 yields the best result.

\paragraph{Independently RNN}
We complete our empirical evaluations of PCFF with an RNN test case. For this we used IndRNNs with 6 layers to solve the sequential/permuted MNIST classification problems based on an existing Implementation\footnote{\url{https://github.com/Sunnydreamrain/IndRNN_pytorch}} provided by the authors of IndRNN \citep{li2018indrnn,li2019indrnn}. We tested over dropout with drop rate 0.1 and PCFF with sample rate 0.1 and report the average test accuracy of three runs. We notice that, while in the permuted MNIST case both dropout (0.9203) and PCFF (0.9145) improves the result (0.9045), in the sequential MNIST case, dropout (0.9830) seems to worsen the performance (0.9841) whereas PCFF (0.9842) delivers comparable result.

\section{Conclusions and discussions}

In this work, we show that neural networks can be interpreted as layered chain graphs, and that feed-forward can be viewed as an approximate inference procedure for these models. This chain graph interpretation provides a unified theoretical framework that elucidates the underlying mechanism of real-world neural networks and provides coherent and in-depth theoretical support for a wide range of empirically established network designs. Furthermore, it also offers a solid foundation to derive new deep learning approaches, with the additional help from the rich existing work on PGMs. It is thus a promising alternative neural network interpretation that deepens our theoretical understanding and unveils a new perspective for future deep learning research.

In the future, we plan to investigate a number of open questions that stem from this work, especially:

\begin{itemize}
\item Is the current chain graph interpretation sufficient to capture the full essence of neural networks? Based on the current results, we are reasonably optimistic that the proposed interpretation can cover an essential part of the neural network mechanism. However, compared to the function approximator view, it only covers a subset of existing techniques. Is this subset good enough?
\item On a related note: can we find chain graph interpretations for other important network designs (or otherwise some chain graph interpretable alternatives with comparable or better performance)? The current work provides a good start, but it is by no means an exhaustive study.
\item Finally, what other new deep learning models and procedures can we build up based on the chain graph framework? The partially collapsed feed-forward inference proposed in this work is just a simple illustrative example, and we believe that many other promising deep learning techniques can be derived from the proposed chain graph interpretation.
\end{itemize}