# How to express a fully connected neural network succintly using linear algebra?

I'm currently reading the paper Federated Learning with Matched Averaging (2020), where the authors claim:

A basic fully connected (FC) NN can be formulated as: $$\hat{y} = \sigma(xW_1)W_2$$ [...]

Expanding the preceding expression $$\hat{y} = \sum_{i=1}^{L} W_{2, i \cdot } \sigma(\langle x, W_{1,\cdot i} \rangle))$$, where $$i\cdot$$ and $$\cdot i$$ denote the ith row and column correspondingly and $$L$$ is the number of hidden units.

I'm having a hard time wrapping my head around how it can be boiled down to this. Is this rigorous? Specifically, what is meant by the ith row and column? Is this formula for only one layer or does it work with multiple layers?

The equation $$\hat{y} = \sigma(xW_\color{green}{1})W_\color{blue}{2} \tag{1}\label{1}$$ is the equation of the forward pass of a single-hidden layer fully connected and feedforward neural network, i.e. a neural network with 3 layers, 1 input layer, 1 hidden layer, and 1 output layer, where

• the input layer is connected to the hidden layer (all scalar inputs are connected to every neuron in the hidden layer)
• the hidden layer is connected to the output layer (all neurons in the hidden layer are connected to all neurons in the output layer

As an example, suppose that we have $$N$$ real-valued features, there are $$L$$ hidden units (or neurons) and $$M$$ output units, then the elements (feature vector and parameters) of equation \ref{1} would have the following shape

• $$x \in \mathbb{R}^{1 \times N}$$
• $$W_\color{green}{1} \in \mathbb{R}^{N \times L}$$
• $$W_\color{blue}{2} \in \mathbb{R}^{L \times M}$$

$$\sigma$$ is an activation function that is applicable to all elements of the matrix separately (i.e. component-wise). So, $$\hat{y} \in \mathbb{R}^{1 \times M}$$.

The equation

$$\hat{y} = \sum_{\color{red}{i} = 1}^{L} W_{\color{blue}{2}, \color{red}{i} \cdot} \sigma(\langle x, W_{\color{green}{1}, \cdot \color{red}{i}} \rangle )\label{2}\tag{2}$$

is another way of writing equation \ref{1}.

Before going to the explanation, let's try to understand equation \ref{2} and its components.

• $$W_{\color{green}{1}, \cdot \color{red}{i}} \in \mathbb{R}^N$$ is the $$\color{red}{i}$$th column of the matrix that connects the inputs to the hidden neurons, so it is a vector of $$N$$ elements (note that we sum over the number of hidden neurons, $$L$$).

• Similarly, $$W_{\color{blue}{2}, \color{red}{i} \cdot} \in \mathbb{R}^M$$ is also a vector, but, in this case, it is a row of the matrix $$W_{\color{blue}{2}}$$ (rather than a column: why? because we use $$\color{red}{i} \cdot$$ instead of $$\cdot \color{red}{i}$$, which refers to the column).

• So, $$\langle x, W_{\color{green}{1}, \cdot \color{red}{i}} \rangle$$ is the dot (or scalar) product between the feature vector $$x$$ and the $$\color{red}{i}$$th column of the matrix that connects the inputs to the hidden neurons, so it's a number (or scalar). Note that both $$x$$ and $$W_{\color{green}{1}, \cdot \color{red}{i}}$$ have $$N$$ elements, so the dot product is well-defined in this case.

• $$W_{\color{blue}{2}, \color{red}{i}} \sigma(\langle x, W_{\color{green}{1}, \color{red}{i}} \rangle))$$ is the product between a vector of shape $$M$$ and a number $$\sigma(\langle x, W_{\color{green}{1}, \cdot \color{red}{i}} \rangle )$$. This is also well-defined. You can multiply a real-number with a vector, it's like multiplying the real-number with each element of the vector.

• $$\sum_{\color{red}{i} = 1}^{L} W_{\color{blue}{2}, \color{red}{i} \cdot} \sigma(\langle x, W_{\color{green}{1}, \cdot \color{red}{i}} \rangle )$$ is thus the sum of $$L$$ vectors of size $$M$$, which makes $$\hat{y}$$ also have size $$M$$, as in equation \ref{1}.

Now, the question is: is equation \ref{2} really equivalent to equation \ref{1}? This is still not easy to see because $$xW_\color{green}{1}$$ is a vector of shape $$L$$, but, in equation \ref{2}, we do not have any vector of shape $$L$$, but we have vectors of shape $$N$$ and $$M$$ (and the vectors of shape $$M$$ are summed $$L$$ times). First, note that $$\sigma(xW_\color{green}{1}) = h\in \mathbb{R}^L$$, so $$hW_\color{blue}{2}$$ are $$M$$ dot products collected in a vector (i.e. $$\hat{y} \in \mathbb{R}^M$$), where the $$j$$th element of $$\hat{y}$$ was computed as a summation of $$L$$ elements (a dot product of two vectors is the element-wise multiplication of the elements of the vectors followed by the summation of these multiplications). Ha, still not clear!

The easiest way (for me) to see that they are equivalent is to think that $$\sigma(xW_\color{green}{1})$$ is a vector of $$L$$ elements $$\sigma(xW_\color{green}{1}) = \ell = [l_1, l_2, \dots, l_L]$$. Then you know that to multiply this vector with $$\ell$$ (from the left), you actually perform a dot product between $$\ell$$ and each column $$W_\color{blue}{2}$$. A dot product is essentially a sum, and that's why we sum in equation \ref{2}. So, essentially, in equation \ref{2}, we first multiple $$l_1$$ with the first row of $$W_\color{blue}{2}$$ (i.e. by all elements of the first row). Then we multiply $$l_2$$ by the second row of $$W_\color{blue}{2}$$. We do this for all $$L$$ rows, then we sum the rows (to conclude the dot product). So, you can think of equation 2 as first perform all multiplications, then summing, rather than dot product-by-dot product.

So, in my head, I have the following picture. To simplify the notation, let $$A$$ denote $$W_\color{blue}{2}$$, so $$A_{ij}$$ is the element at the $$i$$th row and $$j$$th column of matrix $$W_\color{blue}{2}$$. So, we have the following initial matrix

$$A = \begin{bmatrix} A_{11} & A_{12} & A_{13} & \dots & A_{1M} \\ A_{21} & A_{22} & A_{23} & \dots & A_{2M} \\ \vdots & \vdots & \vdots & \dots & \vdots \\ A_{L1} & A_{L2} & A_{L3} & \dots & A_{LM} \\ \end{bmatrix} = \begin{bmatrix} W_{\color{blue}{2}, \color{red}{1} \cdot} \\ W_{\color{blue}{2}, \color{red}{2} \cdot} \\ \vdots \\ W_\color{blue}{2, \color{red}{L} \cdot} \\ \end{bmatrix}$$

Then, in the first iteration of equation \ref{2}, we do the following

$$\begin{bmatrix} l_1 A_{11} & l_1 A_{12} & l_1 A_{13} & \dots & l_1 A_{1M} \\ A_{21} & A_{22} & A_{23} & \dots & A_{2M} \\ \vdots & \vdots & \vdots & \dots & \vdots \\ A_{L1} & A_{L2} & A_{L3} & \dots & A_{LM} \\ \end{bmatrix}$$

In the second, we do the following

$$\begin{bmatrix} l_1 A_{11} & l_1 A_{12} & l_1 A_{13} & \dots & l_1 A_{1M} \\ l_2 A_{21} & l_2 A_{22} & l_2 A_{23} & \dots & l_2 A_{2M} \\ \vdots & \vdots & \vdots & \dots & \vdots \\ A_{L1} & A_{L2} & A_{L3} & \dots & A_{LM} \\ \end{bmatrix}$$ Until we have

$$\begin{bmatrix} l_1 A_{11} & l_1 A_{12} & l_1 A_{13} & \dots & l_1 A_{1M} \\ l_2 A_{21} & l_2 A_{22} & l_2 A_{23} & \dots & l_2 A_{2M} \\ \vdots & \vdots & \vdots & \dots & \vdots \\ l_L A_{L1} & l_L A_{L2} & l_L A_{L3} & \dots & l_L A_{LM} \\ \end{bmatrix}$$ Then we do a reduce sum across the rows to end the dot product (i.e. for each column we sum the elements in the rows). This is exactly equivalent to first performing the dot product between $$\ell$$ and the first column of $$W_\color{blue}{2}$$, then the second column, and so on.