Path-Normalized Optimization of Recurrent Neural Networks with ReLU Activations

Introduction

Recurrent Neural Networks (RNNs) have been found to be successful in a variety of sequence learning problems , including those involving long term dependencies (e.g., ). However, most of the empirical success has not been with “plain” RNNs but rather with alternate, more complex structures, such as Long Short-Term Memory (LSTM) networks or Gated Recurrent Units (GRUs) . Much of the motivation for these more complex models is not so much because of their modeling richness, but perhaps more because they seem to be easier to optimize. As we discuss in Section 3, training plain RNNs using gradient-descent variants seems problematic, and the choice of the activation function could cause a problem of vanishing gradients or of exploding gradients.

In this paper our goal is to better understand the geometry of plain RNNs, and develop better optimization methods, adapted to this geometry, that directly learn plain RNNs with ReLU activations. One motivation for insisting on plain RNNs, as opposed to LSTMs or GRUs, is because they are simpler and might be more appropriate for applications that require low-complexity design such as in mobile computing platforms . In other applications, it might be better to solve optimization issues by better optimization methods rather than reverting to more complex models. Better understanding optimization of plain RNNs can also assist us in designing, optimizing and intelligently using more complex RNN extensions.

Improving training RNNs with ReLU activations has been the subject of some recent attention, with most research focusing on different initialization strategies . While initialization can certainly have a strong effect on the success of the method, it generally can at most delay the problem of gradient explosion during optimization. In this paper we take a different approach that can be combined with any initialization choice, and focus on the dynamics of the optimization itself.

Any local search method is inherently tied to some notion of geometry over the search space (e.g. the space of RNNs). For example, gradient descent (including stochastic gradient descent) is tied to the Euclidean geometry and can be viewed as steepest descent with respect to the Euclidean norm. Changing the norm (even to a different quadratic norm, e.g. by representing the weights with respect to a different basis in parameter space) results in different optimization dynamics. We build on prior work on the geometry and optimization in feed-forward networks, which uses the path-norm (defined in Section 4) to determine a geometry leading to the path-SGD optimization method. To do so, we investigate the geometry of RNNs as feedforward networks with shared weights (Section 2) and extend a line of work on Path-Normalized optimization to include networks with shared weights. We show that the resulting algorithm (Section 4) has similar invariance properties on RNNs as those of standard path-SGD on feedforward networks, and can result in better optimization with less sensitivity to the scale of the weights.

Recurrent Neural Nets as Feedforward Nets with Shared Weights

We view Recurrent Neural Networks (RNNs) as feedforward networks with shared weights.

If the mapping $\pi$ is a one-to-one mapping, then there is no weight sharing and it corresponds to standard feedforward networks. On the other hand, weight sharing exists if $\pi$ is a many-to-one mapping. Two well-known examples of feedforward networks with shared weights are convolutional and recurrent networks. We mostly use the general notation of feedforward networks with shared weights throughout the paper as this will be more general and simplifies the development and notation. However, when focusing on RNNs, it is helpful to discuss them using a more familiar notation which we briefly introduce next.

Time-unfolded RNNs are feedforward networks with shared weights that map an input sequence to an output sequence. Each input node corresponds to either a coordinate of the input vector at a particular time step or a hidden unit at time . Each output node also corresponds to a coordinate of the output at a specific time step. Finally, each internal node refers to some hidden unit at time $t\geq 1$. When discussing RNNs, it is useful to refer to different layers and the values calculated at different time-steps. We use a notation for RNN structures in which the nodes are partitioned into layers and $\mathbf{h}_{t}^{i}$ denotes the output of nodes in layer $i$ at time step $t$. Let $\mathbf{x}=(\mathbf{x}_{1},\dots,\mathbf{x}_{T})$ be the input at different time steps where $T$ is the maximum number of propagations through time and we refer to it as the length of the RNN. For $0\leq i<d$, let $\mathbf{W}_{\text{in}}^{i}$ and $\mathbf{W}_{\text{rec}}^{i}$ be the input and recurrent parameter matrices of layer $i$ and $\mathbf{W}_{\text{out}}$ be the output parameter matrix. Table 1 shows forward computations for RNNs.The output of the function implemented by RNN can then be calculated as $f_{\mathbf{W},t}(x)=h_{t}^{d}$. Note that in this notations, weight matrices $\mathbf{W}_{\text{in}}$, $\mathbf{W}_{\text{rec}}$ and $\mathbf{W}_{\text{out}}$ correspond to “free” parameters of the model that are shared in different time steps.

Non-Saturating Activation Functions

The choice of activation function for neural networks can have a large impact on optimization. We are particularly concerned with the distinction between “saturating” and “non-starting” activation functions. We consider only monotone activation functions and say that a function is “saturating” if it is bounded—this includes, e.g. sigmoid, hyperbolic tangent and the piecewise-linear ramp activation functions. Boundedness necessarily implies that the function values converge to finite values at negative and positive infinity, and hence asymptote to horizontal lines on both sides. That is, the derivative of the activation converges to zero as the input goes to both $-\infty$ and $+\infty$. Networks with saturating activations therefore have a major shortcoming: the vanishing gradient problem . The problem here is that the gradient disappears when the magnitude of the input to an activation is large (whether the unit is very “active” or very “inactive”) which makes the optimization very challenging.

While sigmoid and hyperbolic tangent have historically been popular choices for fully connected feedforward and convolutional neural networks, more recent works have shown undeniable advantages of non-saturating activations such as ReLU, which is now the standard choice for fully connected and Convolutional networks . Non-saturating activations, including the ReLU, are typically still bounded from below and asymptote to a horizontal line, with a vanishing derivative, at $-\infty$. But they are unbounded from above, enabling their derivative to remain bounded away from zero as the input goes to $+\infty$. Using ReLUs enables gradients to not vanish along activated paths and thus can provide a stronger signal for training.

However, for recurrent neural networks, using ReLU activations is challenging in a different way, as even a small change in the direction of the leading eigenvector of the recurrent weights could get amplified and potentially lead to the explosion in forward or backward propagation .

To understand this, consider a long path from an input in the first element of the sequence to an output of the last element, which passes through the same RNN edge at each step (i.e. through many edges in some $E_{i}$ in the shared-parameter representation). The length of this path, and the number of times it passes through edges associated with a single parameter, is proportional to the sequence length, which could easily be a few hundred or more. The effect of this parameter on the path is therefore exponentiated by the sequence length, as are gradient updates for this parameter, which could lead to parameter explosion unless an extremely small step size is used.

Understanding the geometry of RNNs with ReLUs could helps us deal with the above issues more effectively. We next investigate some properties of geometry of RNNs with ReLU activations.

Feedforward networks (with or without shared weights) are highly over-parameterized, i.e. there are many parameter settings ${\mathbf{p}}$ that represent the same function $f_{\mathbf{p}}$. Since our true object of interest is the function $f$, and not the identity ${\mathbf{p}}$ of the parameters, it would be beneficial if optimization would depend only on $f_{\mathbf{p}}$ and not get “distracted” by difference in ${\mathbf{p}}$ that does not affect $f_{\mathbf{p}}$. It is therefore helpful to study the transformations on the parameters that will not change the function presented by the network and come up with methods that their performance is not affected by such transformations.

We say a network $\mathcal{N}$ is invariant to a transformation $\mathcal{T}$ if for any parameter setting $\mathbf{p}$, $f_{\mathbf{p}}=f_{\mathcal{T}(\mathbf{p})}$. Similarly, we say an update rule $\mathcal{A}$ is invariant to $\mathcal{T}$ if for any $\mathbf{p}$, $f_{\mathcal{A}(\mathbf{p})}=f_{\mathcal{A}(\mathcal{T}(\mathbf{p}))}$.

Invariances have also been studied as different mappings from the parameter space to the same function space while we define the transformation as a mapping inside a fixed parameter space. A very important invariance in feedforward networks is node-wise rescaling . For any internal node $v$ and any scalar $\alpha>0$, we can multiply all incoming weights into $v$ (i.e. $w_{u\rightarrow v}$ for any $(u\rightarrow v)\in E$) by $\alpha$ and all the outgoing weights (i.e. $w_{v\rightarrow u}$ for any $(v\rightarrow u)\in E$) by $1/\alpha$ without changing the function computed by the network. Not all node-wise rescaling transformations can be applied in feedforward nets with shared weights. This is due to the fact that some weights are forced to be equal and therefore, we are only allowed to change them by the same scaling factor.

Given a network $\mathcal{N}$, we say an invariant transformation $\widetilde{\mathcal{T}}$ that is defined over edge weights (rather than parameters) is feasible for parameter mapping $\pi$ if the shared weights remain equal after the transformation, i.e. for any $i$ and for any $e,e^{\prime}\in E_{i}$, $\widetilde{\mathcal{T}}(\mathbf{w})_{e}=\widetilde{\mathcal{T}}(\mathbf{w})_{e^{\prime}}$.

Therefore, it is helpful to understand what are the feasible node-wise rescalings for RNNs. In the following theorem, we characterize all feasible node-wise invariances in RNNs.

For any $\alpha$ such that $\alpha^{i}_{j}>0$, any Recurrent Neural Network with ReLU activation is invariant to the transformation $\mathcal{T}_{\alpha}\left(\left[\mathbf{W}_{\text{in}},\mathbf{W}_{\text{rec}},\mathbf{W}_{\text{out}}\right]\right)=\left[\mathcal{T}_{\text{in},\alpha}\left(\mathbf{W}_{\text{in}}\right),\mathcal{T}_{\text{rec},\alpha}\left(\mathbf{W}_{\text{rec}}\right),\mathcal{T}_{\text{out},\alpha}\left(\mathbf{W}_{\text{out}}\right)\right]$ where for any $i,j,k$:

Furthermore, any feasible node-wise rescaling transformation can be presented in the above form.

The proofs of all theorems and lemmas are given in Appendix A. The above theorem shows that there are many transformations under which RNNs represent the same function. An example of such invariances is shown in Fig. 1. Therefore, we would like to have optimization algorithms that are invariant to these transformations and in order to do so, we need to look at measures that are invariant to such mappings.

Path-SGD for Networks with Shared Weights

As we discussed, optimization is inherently tied to a choice of geometry, here represented by a choice of complexity measure or “norm”The path-norm which we define is a norm on functions, not on weights, but as we prefer not getting into this technical discussion here, we use the term “norm” very loosely to indicate some measure of magnitude .. Furthermore, we prefer using an invariant measure which could then lead to an invariant optimization method. In Section 4.1 we introduce the path-regularizer and in Section 4.2, the derived Path-SGD optimization algorithm for standard feed-forward networks. Then in Section 4.3 we extend these notions also to networks with shared weights, including RNNs, and present two invariant optimization algorithms based on it. In Section 4.4 we show how these can be implemented efficiently using forward and backward propagations.

The path-regularizer is the sum over all paths from input nodes to output nodes of the product of squared weights along the path. To define it formally, let $\mathcal{P}$ be the set of directed paths from input to output units so that for any path $\zeta=\left(\zeta_{0},\dots,\zeta_{\text{len}(\zeta)}\right)\in\mathcal{P}$ of length $\text{len}(\zeta)$, we have that $\zeta_{0}\in V_{\text{in}}$, $\zeta_{\text{len}(\zeta)}\in V_{\text{out}}$ and for any $0\leq i\leq\text{len}(\zeta)-1$, $(\zeta_{i}\rightarrow\zeta_{i+1})\in E$. We also abuse the notation and denote $e\in\zeta$ if for some $i$, $e=(\zeta_{i},\zeta_{i+1})$. Then the path regularizer can be written as:

Equivalently, the path-regularizer can be defined recursively on the nodes of the network as:

2 Path-SGD for Feedforward Networks

Path-SGD is an approximate steepest descent step with respect to the path-norm. More formally, for a network without shared weights, where the parameters are the weights themselves, consider the diagonal quadratic approximation of the path-regularizer about the current iterate $\mathbf{w}^{(t)}$:

Using the corresponding quadratic norm $\left\lVert{\mathbf{w}-\mathbf{w}^{\prime}}\right\rVert_{\hat{\gamma}^{2}_{\rm net}(\mathbf{w}^{(t)}+\Delta\mathbf{w})}^{2}=\frac{1}{2}\sum_{e\in E}\frac{\partial^{2}\gamma^{2}_{\rm net}}{\partial w^{2}_{e}}\left(w_{e}-w^{\prime}_{e}\right)^{2}$, we can define an approximate steepest descent step as:

The stochastic version that uses a subset of training examples to estimate $\frac{\partial L}{\partial w_{u\rightarrow v}}(\mathbf{w}^{(t)})$ is called Path-SGD . We now show how Path-SGD can be extended to networks with shared weights.

3 Extending to Networks with Shared Weights

When the networks has shared weights, the path-regularizer is a function of parameters $\mathbf{p}$ and therefore the quadratic approximation should also be with respect to the iterate $\mathbf{p}^{(t)}$ instead of $\mathbf{w}^{(t)}$ which results in the following update rule:

where $\left\lVert{\mathbf{p}-\mathbf{p}^{\prime}}\right\rVert_{\hat{\gamma}^{2}_{\rm net}(\mathbf{p}^{(t)}+\Delta\mathbf{p})}^{2}=\frac{1}{2}\sum_{i=1}^{m}\frac{\partial^{2}\gamma^{2}_{\rm net}}{\partial p^{2}_{i}}\left(p_{i}-p^{\prime}_{i}\right)^{2}$. Solving (7) gives the following update:

The second derivative terms $\kappa_{i}$ are specified in terms of their path structure as follows:

$\kappa_{i}(\mathbf{p})=\kappa^{(1)}_{i}(\mathbf{p})+\kappa^{(2)}_{i}(\mathbf{p})$ where

and $\kappa_{e}(\mathbf{w})$ is defined in (6).

The second term $\kappa^{(2)}_{i}(\mathbf{p})$ measures the effect of interactions between edges corresponding to the same parameter (edges from the same $E_{i}$) on the same path from input to output. In particular, if for any path from an input unit to an output unit, no two edges along the path share the same parameter, then $\kappa^{(2)}(\mathbf{p})=0$. For example, for any feedforward or Convolutional neural network, $\kappa^{(2)}(\mathbf{p})=0$. But for RNNs, there certainly are multiple edges sharing a single parameter on the same path, and so we could have $\kappa^{(2)}(\mathbf{p})\neq 0$.

The above lemma gives us a precise update rule for the approximate steepest descent with respect to the path-regularizer. The following theorem confirms that the steepest descent with respect to this regularizer is also invariant to all feasible node-wise rescaling for networks with shared weights.

For any feedforward networks with shared weights, the update (8) is invariant to all feasible node-wise rescalings. Moreover, a simpler update rule that only uses $\kappa^{(1)}_{i}(\mathbf{p})$ in place of $\kappa_{i}(\mathbf{p})$ is also invariant to all feasible node-wise rescalings.

Equations (9) and (10) involve a sum over all paths in the network which is exponential in depth of the network. However, we next show that both of these equations can be calculated efficiently.

4 Simple and Efficient Computations for RNNs

We show how to calculate $\kappa^{(1)}_{i}(\mathbf{p})$ and $\kappa^{(2)}_{i}(\mathbf{p})$ by considering a network with the same architecture but with squared weights:

Experiments

As we discussed in section 4.4, the second term $\kappa^{(2)}$ in the update rule can be computationally expensive for large networks. In this section we investigate the significance of the second term and show that at least in our experiments, the contribution of the second term is negligible. To compare the two terms $\kappa^{(1)}$ and $\kappa^{(2)}$, we train a single layer RNN with $H=200$ hidden units for the task of word-level language modeling on Penn Treebank (PTB) Corpus . Fig. 2 compares the performance of SGD vs. Path-SGD with/without $\kappa^{(2)}$. We clearly see that both version of Path-SGD are performing very similarly and both of them outperform SGD significantly. This results in Fig. 2 suggest that the first term is more significant and therefore we can ignore the second term.

To better understand the importance of the two terms, we compared the ratio of the norms $\left\lVert{\kappa^{(2)}}\right\rVert_{2}/\left\lVert{\kappa^{(1)}}\right\rVert_{2}$ for different RNN lengths $T$ and number of hidden units $H$. The table in Fig. 2 shows that the contribution of the second term is bigger when the network has fewer number of hidden units and the length of the RNN is larger ($H$ is small and $T$ is large). However, in many cases, it appears that the first term has a much bigger contribution in the update step and hence the second term can be safely ignored. Therefore, in the rest of our experiments, we calculate the Path-SGD updates only using the first term $\kappa^{(1)}$.

2 Synthetic Problems with Long-term Dependencies

Training Recurrent Neural Networks is known to be hard for modeling long-term dependencies due to the gradient vanishing/exploding problem . In this section, we consider synthetic problems that are specifically designed to test the ability of a model to capture the long-term dependency structure. Specifically, we consider the addition problem and the sequential MNIST problem.

Addition problem: The addition problem was introduced in . Here, each input consists of two sequences of length $T$, one of which includes numbers sampled from the uniform distribution with range $$and the other sequence serves as a mask which is filled with zeros except for two entries. These two entries indicate which of the two numbers in the first sequence we need to add and the task is to output the result of this addition. Sequential MNIST: In sequential MNIST, each digit image is reshaped into a sequence of length$784$, turning the digit classification task into sequence classification with long-term dependencies .

For both tasks, we closely follow the experimental protocol in . We train a single-layer RNN consisting of 100 hidden units with path-SGD, referred to as RNN-Path. We also train an RNN of the same size with identity initialization, as was proposed in , using SGD as our baseline model, referred to as IRNN. We performed grid search for the learning rates over $\{10^{-2},10^{-3},10^{-4}\}$ for both our model and the baseline. Non-recurrent weights were initialized from the uniform distribution with range $[-0.01,0.01]$. Similar to , we found the IRNN to be fairly unstable (with SGD optimization typically diverging). Therefore for IRNN, we ran 10 different initializations and picked the one that did not explode to show its performance.

In our first experiment, we evaluate Path-SGD on the addition problem. The results are shown in Fig. 3 with increasing the length $T$ of the sequence: $\{100,400,750\}$. We note that this problem becomes much harder as $T$ increases because the dependency between the output (the sum of two numbers) and the corresponding inputs becomes more distant. We also compare RNN-Path with the previously published results, including identity initialized RNN (IRNN), unitary RNN (uRNN), and np-RNNThe original paper does not include any result for 750, so we implemented np-RNN for comparison. However, in our implementation the np-RNN is not able to even learn sequences of length of 200. Thus we put “¿2” for length of 750. introduced by . Table 3 shows the effectiveness of using Path-SGD. Perhaps more surprisingly, with the help of path-normalization, a simple RNN with the identity initialization is able to achieve a 0% error on the sequences of length 750, whereas all the other methods, including LSTMs, fail. This shows that Path-SGD may help stabilize the training and alleviate the gradient problem, so as to perform well on longer sequence. We next tried to model the sequences length of 1000, but we found that for such very long sequences RNNs, even with Path-SGD, fail to learn.

Next, we evaluate Path-SGD on the Sequential MNIST problem. Table 3, right column, reports test error rates achieved by RNN-Path compared to the previously published results. Clearly, using Path-SGD helps RNNs achieve better generalization. In many cases, RNN-Path outperforms other RNN methods (except for LSTMs), even for such a long-term dependency problem.

3 Language Modeling Tasks

In this section we evaluate Path-SGD on a language modeling task. We consider two datasets, Penn Treebank (PTB-c) and text8 http://mattmahoney.net/dc/textdata. PTB-c: We performed experiments on a tokenized Penn Treebank Corpus, following the experimental protocol of . The training, validations and test data contain 5017k, 393k and 442k characters respectively. The alphabet size is 50, and each training sequence is of length 50. text8: The text8 dataset contains 100M characters from Wikipedia with an alphabet size of 27. We follow the data partition of , where each training sequence has a length of 180. Performance is evaluated using bits-per-character (BPC) metric, which is $\log_{2}$ of perplexity.

Similar to the experiments on the synthetic datasets, for both tasks, we train a single-layer RNN consisting of 2048 hidden units with path-SGD (RNN-Path). Due to the large dimension of hidden space, SGD can take a fairly long time to converge. Instead, we use Adam optimizer to help speed up the training, where we simply use the path-SGD gradient as input to the Adam optimizer.

We also train three additional baseline models: a ReLU RNN with 2048 hidden units, a tanh RNN with 2048 hidden units, and an LSTM with 1024 hidden units, all trained using Adam. We performed grid search for learning rate over $\{10^{-3},5\cdot 10^{-4},10^{-4}\}$ for all of our models. For ReLU RNNs, we initialize the recurrent matrices from uniform$[-0.01,0.01]$, and uniform$[-0.2,0.2]$ for non-recurrent weights. For LSTMs, we use orthogonal initialization for the recurrent matrices and uniform$[-0.01,0.01]$ for non-recurrent weights. The results are summarized in Table 3.

We also compare our results to an RNN that uses hidden activation regularizer (TRec,$\beta=500$), Multiplicative RNNs trained by Hessian Free methods (HF-MRNN), and an RNN with smooth version of ReLU . Table 3 shows that path-normalization is able to outperform RNN-ReLU and RNN-tanh, while at the same time shortening the performance gap between plain RNN and other more complicated models (e.g. LSTM by 57% on PTB and 54% on text8 datasets). This demonstrates the efficacy of path-normalized optimization for training RNNs with ReLU activation.

Conclusion

We investigated the geometry of RNNs in a broader class of feedforward networks with shared weights and showed how understanding the geometry can lead to significant improvements on different learning tasks. Designing an optimization algorithm with a geometry that is well-suited for RNNs, we closed over half of the performance gap between vanilla RNNs and LSTMs. This is particularly useful for applications in which we seek compressed models with fast prediction time that requires minimum storage; and also a step toward bridging the gap between LSTMs and RNNs.

This research was supported in part by an NSF RI-AF award and by Intel ICRI-CI. We thank Saizheng Zhang for sharing a base code for RNNs.

References

Appendix A Proofs

We first show that any RNN is invariant to $\mathcal{T}_{\alpha}$ by induction on layers and time-steps. More specifically, we prove that for any $0\leq t\leq T$ and $1\leq i<d$, $\mathbf{h}_{t}^{i}\left(\mathcal{T}_{\alpha}(\mathbf{W})\right)[j]=\alpha^{i}_{j}\mathbf{h}_{t}^{i}(\mathbf{W})[j]$. The statement is clearly true for $t=0$; because for any $i,j$, $\mathbf{h}_{0}^{i}\left(\mathcal{T}_{\alpha}(\mathbf{W})\right)[j]=\alpha^{i}_{j}\mathbf{h}_{0}^{i}(\mathbf{W})[j]=0$.

Next, we show that for $i=1$, if we assume that the statement is true for $t=t^{\prime}$, then it is also true for $t=t^{\prime}+1$:

We now need to prove the statement for $1<i<d$. Assuming that the statement is true for $t\leq t^{\prime}$ and the layers before $i$, we have:

Finally, we can show that the output is invariant for any $j$ at any time step $t$:

We now show that any feasible node-wise rescaling can be presented as $\mathcal{T}_{\alpha}$. Recall that node-wise rescaling invariances for a general feedforward network can be written as $\widetilde{\mathcal{T}_{\beta}}(\mathbf{w})_{u\rightarrow v}=(\beta_{v}/\beta_{u})w_{u\rightarrow v}$ for some $\beta$ where $\beta_{v}>0$ for internal nodes and $\beta_{v}=1$ for any input/output nodes. An RNN with $T=0$ has no weight sharing and for each node $v$ with index $j$ in layer $i$, we have $\beta_{v}=\alpha_{j}^{i}$. For any $T>0$ however, we there is no invariance that is not already counted. The reason is that by fixing the values of $\beta_{v}$ for the nodes in time step 0, due to the feasibility, the values of $\beta$ for nodes in other time-steps should be tied to the corresponding value in time step . Therefore, all invariances are included and can be presented in form of $\mathcal{T}_{\alpha}$.

A.2 Proof of Lemma 1

We prove the statement simply by calculating the second derivative of the path-regularizer with respect to each parameter:

Taking the second derivative then gives us both terms after a few calculations:

A.3 Proof of Theorem 2

Node-wise rescaling invariances for a feedforward network can be written as $\mathcal{T}_{\beta}(\mathbf{w})_{u\rightarrow v}=(\beta_{v}/\beta_{u})w_{u\rightarrow v}$ for some $\beta$ where $\beta_{v}>0$ for internal nodes and $\beta_{v}=1$ for any input/output nodes. Any feasible invariance for a network with shared weights can also be written in the same form. The only difference is that some of $\beta_{v}$s are now tied to each other in a way that shared weights have the same value after transformation. First, note that since the network is invariant to the transformation, the following statement holds by an induction similar to Theorem 1 but in the backward direction:

for any $(u\rightarrow v)\in E$. Furthermore, by the proof of the Theorem 1 we have that for any $(u\rightarrow v)\in E$, $h_{u}(\mathcal{T}_{\beta}(\mathbf{p}))=\beta_{u}h_{u}(\mathbf{p})$. Therefore,

where $(u^{\prime}\rightarrow v^{\prime})\in E_{i}$. In order to prove the theorem statement, it is enough to show that for any edge $(u\rightarrow v)\in E_{i}$, $\kappa_{i}(\mathcal{T}_{\beta}(\mathbf{p}))=(\beta_{u}/\beta_{v})^{2}\kappa_{i}(\mathbf{p})$ because this property gives us the following update:

Therefore, it is remained to show that for any edge $(u\rightarrow v)\in E_{i}$ $v$, $\kappa_{i}(\mathcal{T}_{\beta}(\mathbf{p}))=(\beta_{u}/\beta_{v})^{2}\kappa_{i}(\mathbf{p})$. We show that this is indeed true for both terms $\kappa^{(1)}$ and $\kappa^{(2)}$ separately.

We first prove the statement for $\kappa^{(1)}$. Consider each path $\zeta\in\mathcal{P}$. By an inductive argument along the path, it is easy to see that multiplying squared weights along this path is invariant to the transformation:

Therefore, we have that for any edge $e\in E$ and any $\zeta\in\mathcal{P}$,

Taking sum over all paths $\zeta\in\mathcal{P}$ and all edges $e=(u\rightarrow v)\in E$ completes the proof for $\kappa^{(1)}$. Similarly for $\kappa^{(2)}$, considering any two edges $e_{1}\neq e_{2}$ and any path $\zeta_{\mathcal{P}}$, we have that:

where $(u\rightarrow v)\in E_{i}$. Again, taking sum over all paths $\zeta$ and all edges $e_{1}\neq e_{2}$ proves the statement for $\kappa^{(2)}$ and consequently for $\kappa^{(1)}+\kappa^{(2)}$.

A.4 Proof of Theorem 3

This path can therefore be decomposed into a path from input to edge $e$ and a path from edge $e$ to the output. Therefore, for any edge $e$, we can factor out the number corresponding to the paths that go through $e$ and rewrite $\kappa^{(1)}$ as follows:

where $\mathcal{P}_{\text{in}\rightarrow u}$ is the set of paths from input nodes to node $v$ and $\mathcal{P}_{v\rightarrow\text{out}}$ is defined similarly for the output nodes.

Therefore, $\kappa^{(1)}$ can be written as:

Next, we show how to calculate the second term, i.e. $\kappa^{(2)}$. Each term in $\kappa^{(2)}$ corresponds to a path that goes through two edges. We can decompose such paths and rewrite $\kappa^{(2)}$ similar to the first term:

where $\mathcal{P}_{u\rightarrow v}$ is the set of all directed paths from node $u$ to node $v$.