\paragraph{Setup.}
To illustrate the numerical properties of our proposed collapsed Taylor mode, we consider a two-layered MLP with element-wise $\mathrm{tanh}$ activation $\vphi: \sR^I \to \sR^I$.
The MLP is denoted by $\vf := \vg \circ \vphi \circ \vh$.
The two linear layers are given as $\vh: \sR^D \to \sR^I, \vh(\vx_0) = \mW_1 \vx_0 + \vb_1$ and $\vg: \sR^I \to \sR^C$, $\vg(\vphi_0) = \mW_2 \vphi_0 + \vb_2$, with weights $\mW_1 \in \sR^{I \times D}, \mW_2 \in \sR^{C \times I}$ and bias $\vb_1 \in \sR^I, \vb_2 \in \sR^C$.
Below we compare the computational and storage complexity, as well as stability for evaluating the sum of the second coefficients $\smash{\sum_{r=1}^R \langle \partial^2 \vf(\vx_0), \vv_i^{\otimes 2} \rangle = \sum_{r=1}^R \vg_{2,r}}$ (see \cref{eq:sum-k-directional}) between collapsed and standard Taylor mode. For this toy example we show (i) collapsing uses less operations and (ii) both methods are similarly stable based on our simplified error analysis.

\paragraph{Computational \& storage complexity.}
Both vanilla and collapsed Taylor mode evaluate the function values ($\vh_0, \vphi_0, \vg_0$) and the first derivatives ($\{\vh_{1,r}, \vphi_{1, r}, \vg_{1, r}\}$) by propagating $1 + R$ coefficients at each layer
\begin{equation*}
  \begin{pmatrix*}[l]
    \vh_0
    =
    \mW_1 \vx_0 + \vb_1
    \\
    \left\{
    \vh_{1,r}
    \right\}
    =
    \left\{
    \left\langle \mW_1, \vx_{1,r}\right\rangle
    \right\}
  \end{pmatrix*}
  \overset{\text{(\ref{eq:faa-di-bruno})}}{\to}
  \begin{pmatrix*}[l]
    \vphi_0
    =
    \vphi(\vh_0)
    \\
    \left\{
    \vphi_{1,r}
    \right\}
    =
    \left\{
    \left\langle \partial \vphi(\vh_0), \vh_{1,r} \right\rangle
    \right\}
  \end{pmatrix*}
  \overset{\text{(\ref{eq:faa-di-bruno})}}{\to}
  \begin{pmatrix*}[l]
    \vg_0
    =
    \mW_2 \vphi_0 + \vb_2
    \\
    \left\{
    \vg_{1,r}
    \right\}
    =
    \left\{
    \left\langle\mW_2, \vphi_{1,r}\right\rangle
    \right\}
  \end{pmatrix*}
\end{equation*}
with $\partial \vphi(\vh_0) = \partial \mathrm{tanh}(\vh_0) = \diag(\vone - \vphi_0^{\odot 2}) \in \sR^{I \times I}$ the $\mathrm{tanh}$-activation layer's Jacobian.
The propagation costs $1 + R$ matrix-vector multiplications with $\mW_1$, $1 + R$ matrix-vector multiplications with $\mW_2$, $R$ Hadamard products with the derivative of $\vphi$ (since $\left\langle \diag(\va), \vh_{1,r} \right\rangle = \va \odot \vh_{1,r}$), and one Hadamard product to compute $\partial \vphi(\vh_0)$.
Additionally, there is one vector addition with the bias $\vb_1$, one vector addition with $\vb_2$, one vector subtraction in $\partial \vphi(\vh_0)$ (counted as vector addition), as well as the evaluation of $\vphi(\vh_0)$.
$3+3R$ vectors are stored.

\begin{table}[!b]
  \caption{\textbf{Comparison of theoretical computational and storage complexity} between standard Taylor mode and collapsed Taylor mode for a two-layer MLP computing the sum $\smash{\sum_{r=1}^R \langle \partial^2 \vf(\vx_0), \vv_r^{\otimes 2} \rangle}$.}
  \label{tab:run-time-storage-comparison}
  \vspace{1ex}
  \centering
  \renewcommand{\arraystretch}{1.2}
  \begin{tabular}{l|cc}
    \toprule
    \multicolumn{3}{c}{\textbf{Computational Complexity}} \\
    \midrule
    \textbf{Operation}
    & \textcolor{tab-orange}{Standard Taylor}
    & \textcolor{tab-green}{Collapsed (ours)} \\
    \midrule
    \# Matrix-vector products
    & $4R + 2$
    & $2R + 4$ \\

    \# Hadamard products
    & $4R + 2$
    & $3R + 3$ \\

    \# Vector additions
    & $2R + 2$
    & $2R + 2$ \\

    \# Scalar multiplications
    & $1$
    & $1$ \\

    \# Activation evaluations
    & $I$
    & $I$ \\
    \midrule
    \multicolumn{3}{c}{\textbf{Storage Complexity}} \\
    \midrule
    \# Vectors stored
    & $6R + 3$
    & $3R + 6$ \\
    \bottomrule
  \end{tabular}
\end{table}

For the second derivatives, vanilla Taylor mode propagates $R$ vectors
\begin{equation*}
  \begin{aligned}
    \begin{pmatrix*}[l]
      \left\{\vh_{2,r}\right\}
      =
      \left\{\left\langle \mW_1, \vx_{2,r} \right\rangle\right\}
    \end{pmatrix*}
    &\overset{\text{(\ref{eq:faa-di-bruno})}}{\to}
      \begin{pmatrix*}[l]
        \left\{
        \vphi_{2,r}
        \right\}
        =
        \left\{
        \left\langle \partial^2 \vphi(\vh_0), \vh_{1, r}^{\otimes 2} \right\rangle + \left\langle \partial \vphi(\vh_0), \vh_{2,r} \right \rangle
        \right\}
      \end{pmatrix*}
    \\
    &\overset{\text{(\ref{eq:faa-di-bruno})}}{\to}
      \begin{pmatrix*}[l]
        \left\{
        \vg_{2,r}
        \right\}
        =
        \left\{
        \left\langle\mW_2 , \vphi_{2,r}\right\rangle
        \right\}
      \end{pmatrix*}
  \end{aligned}
\end{equation*}
with activation Hessian $\partial^2 \vphi(\vh_0) \in \sR^{I \times I \times I}$ of entries $[\partial^2 \vphi(\vh_0)]_{i,j,k} = [-2 \vphi_0 \odot (\vone - \vphi_0^{\odot 2})]_i \delta_{i,j,k}$ and contraction $\smash{\left\langle \partial^2 \vphi(\vh_0), \vh_{1, r}^{\otimes 2} \right\rangle = -2 \vphi_0 \odot (\vone - \vphi_0^{\odot 2}) \odot \vh_{1,r}^{\odot 2}}$.
These vectors are summed up to get the result $\smash{\sum_{r=1}^R \langle \partial^2 \vf(\vx_0), \vv_i^{\otimes 2} \rangle = \sum_{r=1}^R \vg_{2,r}}$.
This costs $2R$ matrix-vector products with the weights, $1 + 3R$ Hadamard products, $2R - 1$ vector additions, and a single scalar multiplication.

In contrast, collapsed Taylor mode propagates only a single summed vector
\begin{equation*}
  \begin{aligned}
    \begin{pmatrix*}[l]
      \displaystyle\sum_{r=1}^R \vh_{2,r} = \left\langle \mW_1,  \sum_{r=1}^R\vx_{2,r} \right\rangle
    \end{pmatrix*}
    \overset{\text{(\ref{eq:faa-di-bruno})}}{\to}
    &\begin{pmatrix*}[l]
      \displaystyle\sum_{r=1}^R\vphi_{2,r} = \sum_{r=1}^R \left\langle \partial^2 \vphi(\vh_0),  \vh_{1, r}^{\otimes 2} \right\rangle + \left\langle \partial \vphi(\vh_0), \sum_{r=1}^R \vh_{2,r} \right\rangle
    \end{pmatrix*}
    \\
    \overset{\text{(\ref{eq:faa-di-bruno})}}{\to}
    &\begin{pmatrix*}[l]
      \displaystyle\sum_{r=1}^R\vg_{2,r} = \left\langle\mW_2, \sum_{r=1}^R\vphi_{2,r}\right\rangle
    \end{pmatrix*}.
  \end{aligned}
\end{equation*}
This costs two matrix-vector products, $2 + 2R$ Hadamard products, $2R-1$ vector additions, and a single scalar multiplication. 
\Cref{tab:run-time-storage-comparison} summarizes the accumulated costs.



\paragraph{Error analysis.} For our numerical experiments, the result of all implementations (nested 1\textsuperscript{st}-order and standard/collapsed Taylor mode) was always checked to be close.
To supplement this experimental error analysis, we sketch a simplified error analysis below.
We assume that there are error-prone first- and second-order inputs $\{\vx_{1, r} + \vvarepsilon_{1, r}\}_r$ and $\{\vx_{2, r} + \vvarepsilon_{2, r}\}_r$ with errors $\{\vvarepsilon_{1, r}, \vvarepsilon_{2, r}\}_{r=1}^R$ that can be seen as the error of previous propagation steps.
An error-prone $\vx_0$ would complicate our brief discussion too much and is ignored here.
We consider again $\vf = \vg \circ \vphi \circ \vh$.
The error-influenced coefficients are denoted $\vg_{2, r}^\varepsilon$.

Using vanilla Taylor mode, the erroneous result is
\begin{align*}
  \sum_{r=1}^R \vg^\varepsilon_{2,r}
  &=
    \sum_{r=1}^R \Big( \Big\langle\mW_2,
    \left\langle \partial^2 \vphi(\vh_0),
    \left\langle \mW_1,
    \vx_{1,r} +  \vvarepsilon_{1, r} \right\rangle^{\otimes 2}
    \right\rangle
  +
    \left\langle
    \partial \vphi(\vh_0),
    \left\langle
    \mW_1,
    \vx_{2,r}
    +
    \vvarepsilon_{2,r}
    \right\rangle
    \right\rangle
    \Big\rangle
    \Big)
  \\
  &=
    \sum_{r=1}^R \vg_{2, r}
  \\
  &\phantom{=}+
    \sum_{r=1}^R \Big(\left\langle\mW_2,
    \left\langle \partial^2 \vphi(\vh_0),
    \left\langle \mW_1,
    \vx_{1,r}
    \right\rangle
    \otimes
    \left \langle
    \mW_1,
    \vvarepsilon_{1, r}
    \right \rangle
    \right\rangle
    \right \rangle
    \\
    &\phantom{===}+
    \left\langle\mW_2,
    \left\langle \partial^2 \vphi(\vh_0),
    \left \langle
    \mW_1,
    \vvarepsilon_{1, r}
    \right \rangle
    \otimes
    \left\langle \mW_1,
    \vx_{1,r}
    \right\rangle
    \right\rangle
    \right\rangle
  \\
  &\phantom{===}+
    \left\langle\mW_2,
    \left\langle \partial^2 \vphi(\vh_0),
    \langle \mW_1,
    \vvarepsilon_{1,r}
    \rangle^{\otimes 2}
    \right \rangle
    \right\rangle
    +
    \left\langle\mW_2,
    \left\langle
    \partial \vphi(\vh_0),
    \left\langle
    \mW_1,
    \vvarepsilon_{2,r}
    \right\rangle
    \right\rangle
    \right\rangle
    \Big)
  \\
  &=
    \sum_{r=1}^R \vg_{2, r}
    +
    \Delta \vg_{2,R}^S
    +
    \sum_{r= 1}^R
    \left \langle \mW_2
    \left\langle
    \partial \vphi(\vh_0),
    \left\langle
    \mW_1,
    \vvarepsilon_{2,r}
    \right\rangle
    \right\rangle
    \right\rangle.
\end{align*}
All errors related to the first-order coefficients are summarized in
\begin{align*}
  \Delta \vg_{2,R}^S &:= \sum_{r=1}^R \Big(\left\langle\mW_2,
                       \left\langle \partial^2 \vphi(\vh_0),
                       \left\langle \mW_1,
                       \vx_{1,r}
                       \right\rangle
                       \otimes
                       \left \langle
                       \mW_1,
                       \vvarepsilon_{1, r}
                       \right \rangle
                       \right\rangle
                       \right \rangle
                       \\
                       &\phantom{:==}+
                       \left\langle\mW_2,
                       \left\langle \partial^2 \vphi(\vh_0),
                       \left \langle
                       \mW_1,
                       \vvarepsilon_{1, r}
                       \right \rangle
                       \otimes
                       \left\langle \mW_1,
                       \vx_{1,r}
                       \right\rangle
                       \right\rangle
                       \right\rangle
  \\
                     &\phantom{:==}+
                       \left\langle\mW_2,
                       \left\langle \partial^2 \vphi(\vh_0),
                       \langle \mW_1,
                       \vvarepsilon_{1,r}
                       \rangle^{\otimes 2}
                       \right\rangle
                       \right\rangle
                       \Big)\,.
\end{align*}
The collapsed Taylor mode results in
\begin{align*}
  \sum_{r=1}^R \vg^\varepsilon_{2,r}
  &=
    \left\langle\mW_2,  \sum_{r=1}^R \left(
    \left\langle \partial^2 \vphi(\vh_0),
    \left\langle \mW_1,
    \vx_{1,r} + \vvarepsilon_{1, r} \right\rangle^{\otimes 2}
    \right\rangle
    \right)
    \right\rangle
  \\
  &+
    \left\langle\mW_2,
    \left\langle
    \partial \vphi(\vh_0),
    \left\langle
    \mW_1,
    \sum_{r=1}^R \left(
    \vx_{2,r}
    +
    \vvarepsilon_{2,r}
    \right)
    \right\rangle
    \right\rangle
    \right\rangle
  \end{align*}
  \begin{align*}
  &=
    \sum_{r=1}^R \vg_{2,r}
  \\
  &\phantom{=}+
    \Big\langle
    \mW_2,
    \sum_{r=1}^R
    \Big(
    \langle
    \partial^2 \vphi(\vh_0),
    \langle
    \mW_1, \vx_{1, r}
    \rangle
    \otimes
    \langle
    \mW_1, \vvarepsilon_{1, r}
    \rangle
    \rangle
    +
    \langle
    \partial^2 \vphi(\vh_0),
    \langle
    \mW_1, \vvarepsilon_{1, r}
    \rangle
    \otimes
    \langle
    \mW_1, \vx_{1, r}
    \rangle
    \rangle
  \\
  &\phantom{======}+
    \langle
    \partial^2 \vphi(\vh_0),
    \langle
    \mW_1, \vvarepsilon_{1, r}
    \rangle^{\otimes 2}
    \rangle
    \Big)
    \Big\rangle
  \\
  &\phantom{=}+
    \left \langle
    \mW_2,
    \left \langle \partial \vphi(\vh_0),
    \left\langle \mW_1,
    \sum_{r=1}^R
    \vvarepsilon_{2, r}
    \right\rangle
    \right \rangle
    \right \rangle
  \\
  &=
    \sum_{r=1}^R \vg_{2,r}
    +
    \Delta \vg_{2,R}^C
    +
    \left \langle
    \mW_2,
    \left \langle \partial \vphi(\vh_0),
    \left\langle \mW_1,
    \sum_{r=1}^R
    \vvarepsilon_{2, r}
    \right\rangle
    \right \rangle
    \right \rangle,
\end{align*}
where the first-order errors are collected in
\begin{align*}
  \Delta \vg_{2, R}^C &:=
                        \Big\langle
                        \mW_2,
                        \sum_{r=1}^R
                        \Big(
                        \langle
                        \partial^2 \vphi(\vh_0),
                        \langle
                        \mW_1, \vx_{1, r}
                        \rangle
                        \otimes
                        \langle
                        \mW_1, \vvarepsilon_{1, r}
                        \rangle
                        \rangle
                        \\
                        &\phantom{:=====}+
                        \langle
                        \partial^2 \vphi(\vh_0),
                        \langle
                        \mW_1, \vvarepsilon_{1, r}
                        \rangle
                        \otimes
                        \langle
                        \mW_1, \vx_{1, r}
                        \rangle
                        \rangle
                        +
                        \langle
                        \partial^2 \vphi(\vh_0),
                        \langle
                        \mW_1, \vvarepsilon_{1, r}
                        \rangle^{\otimes 2}
                        \rangle
                        \Big)
                        \Big\rangle\,.
\end{align*}
\paragraph{Error analysis (summary and discussion).} Without considering floating-point operations, the errors are equivalent.
This is not surprising, since our collapsing method is mathematically equivalent to the standard Taylor mode on the same input coefficients.

Incorporating floating-point operations for the function evaluations, inner product, tensor product, and summations would greatly complicate the discussion, which is not part of the paper.
Still, the error could be split into the same three parts for both vanilla and collapsed Taylor mode.
For the first-order errors $\smash{\Delta \vg_{2, R}^S}$ and $\smash{\Delta \vg_{2, R}^C}$, however, even with floating-point operations, the errors are structurally similar, since apart from the most outer inner product (with $\mW_2$) and the summation, all operations are done in the same order.
In practice, we would expect smaller errors for the collapsing method due to the reduced number of operations.
The second error term, which collects the error of the second-order coefficients, could also reduce the accumulation of error terms.
Of course, the actual condition and input, and output dimensions of the matrices are crucial.
Theoretically, this could even lead to a similar error asymptotically.
If inputs are small, one could argue that catastrophic cancellations are more likely to happen in our case, since we sum first.
But note that those cancellations are then also likely to happen in the standard Taylor mode, because weight matrices are often normalized, and the outputs of the activation functions are small if the input is small.

We plan to investigate this more rigorously in the future.

%%% Local Variables:
%%% mode: LaTeX
%%% TeX-master: "../main"
%%% End:
