SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Pieter-Jan Hoedt & Günter Klambauer
LIT AI Lab & ELLIS Unit Linz
Institute for Machine Learning
Johannes Kepler University, Linz, Austria
{hoedt, klambauer}@ml.jku.at
Input-Convex Neural Networks (ICNNs) are networks that guarantee convexity in their input-output mapping. These networks have been successfully applied for energy-based modelling, optimal transport problems and learning invariances. The convexity of ICNNs is achieved by using non-decreasing convex activation functions and non-negative weights. Because of these peculiarities, previous initialisation strategies, which implicitly assume centred weights, are not effective for ICNNs. By studying signal propagation through layers with non-negative weights, we are able to derive a principled weight initialisation for ICNNs. Concretely, we generalise signal propagation theory by removing the assumption that weights are sampled from a centred distribution. In a set of experiments, we demonstrate that our principled initialisation effectively accelerates learning in ICNNs and leads to better generalisation. Moreover, we find that, in contrast to common belief, ICNNs can be trained without skip-connections when initialised correctly. Finally, we apply ICNNs to a real-world drug discovery task and show that they allow for more effective molecular latent space exploration.
Input-Convex Networks. Input-Convex Neural Networks (ICNNs) are networks for which each output neuron is convex with respect to the inputs. The convexity is a result of using non-decreasing convex activation functions and weight matrices with non-negative entries (Amos et al., 2017). ICNNs were originally introduced in the context of energy modelling (Amos et al., 2017). Also in other contexts, ICNNs have proven to be useful. E.g. Sivaprasad et al. (2021) show that ICNNs can be used as regular classification models, Makkuva et al. (2020) rely on the convexity to model optimal transport mappings, and Nesterov et al. (2022) illustrate how ICNNs can be used to learn invariances by simplifying the search for level sets — i.e. inputs for which the output prediction remains unchanged. ICNNs have also been used in model predictive control, where recently an input-convex variant of LSTMs was proposed to reduce convergence time while ensuring closed-loop stability (Wang and Wu, 2023). Despite their successful application for various tasks, convergence can be notably slow (Sivaprasad et al., 2021, Fig. 1 (d)). We hypothesize that this slow training is the result of poor initialisation of the positive weights in ICNNs. This poor initialisation leads to distribution shifts that make learning harder, as illustrated in Figure 1. Therefore, we propose a principled initialisation strategy for layers with non-negative weight matrices (cf. Chang et al., 2020).
Importance of initialisation strategies. Initialisation strategies have enabled faster and more stable learning in deep networks (LeCun et al., 1998; Glorot and Bengio, 2010; He et al., 2015). The goal of a good initialisation strategy is to produce similar statistics in every layer. This can be done in the forward pass (LeCun et al., 1998; Mishkin and Matas, 2016; Klambauer et al., 2017; Chang et al.,
feature correlations
pre-activation distributions
(a) traditional
(b) ICNN
Figure 1: Illustration of the effects due to good signal propagation in hidden layers. Blue and red connections depict negative and positive weights, respectively. The top row shows histograms of pre-activations in each hidden layer. The bottom row displays the feature correlation matrices for these layers. Small visualisations on the left depict the input distribution.
(c) ICNN + init
2020) or during back-propagation (Glorot and Bengio, 2010; Hoedt et al., 2018; Defazio and Bottou, 2021). It is also important to account for the effects due to non-linearities in the network (Saxe et al., 2014; He et al., 2015; Klambauer et al., 2017; Hoedt et al., 2018). However, there are no initialisation strategies for non-negative weight matrices, which are a key component of ICNNs (Amos et al., 2017). We derive an initialisation strategy that accounts for the non-negative weights in ICNNs by generalising the signal propagation principles that underlie modern initialisation strategies.
Signal propagation. The derivation of initialisation strategies typically builds on the signal propagation framework introduced by Neal (1995). This signal propagation theory has been used and expanded in various ways (Saxe et al., 2014; Poole et al., 2016; Klambauer et al., 2017; Martens et al., 2021). One critical assumption in this traditional signal propagation theory is that weights are sampled from a centred distribution, i.e. with zero mean. In ICNNs, this is not possible because some weight matrices are constrained to be non-negative (Amos et al., 2017). Therefore, we generalise the traditional signal propagation theory to allow for non-centred distributions.
Skip-connections in ICNNs. Architectures of ICNNs typically include skip-connections (Amos et al., 2017; Sivaprasad et al., 2021; Makkuva et al., 2020; Nesterov et al., 2022). The skip-connections in ICNNs were introduced to increase their representational power (Amos et al., 2017). Although it is possible to study signal propagation with skip-connections (e.g. Yang and Schoenholz, 2017; Brock et al., 2021; Hoedt et al., 2022), they are typically built on top of existing results for plain networks. Therefore, we limit our theoretical results to ICNNs without skip-connections. However, we find that ICNNs without skip-connections can be successfully trained, indicating that skip-connections might not be necessary for representational power. We show that with our initialisation, we are able to train an ICNN to the same performance as a non-convex baseline without skip-connections. This confirms the hypothesis that good signal propagation can replace skip-connections (Martens et al., 2021; Zhang et al., 2022).
Contributions. Our contributions<sup>1</sup> can be summarised as follows:
- We generalise signal propagation theory to include weights without zero mean (Section 2).
- We derive a principled initialisation strategy for ICNNs from our new theory (Section 3).
- We empirically demonstrate the effectiveness of our initialisation strategy (Section 5.1).
- We apply ICNNs in a real-world drug-discovery setting (Sections 5.2 and 5.3).
2 Generalising Signal Propagation
In this section, we revisit the traditional signal propagation theory (Neal, 1995), which assumes centred weights. This provides us with a framework to study signal propagation in standard fully-connected networks. We then expand this framework to enable the study of networks where weights do not have zero mean to derive a weight initialisation strategy for ICNNs (Section 3).
$^{1}$ Code for figures and experiments can be found at https://github.com/ml-jku/convex-init
2.1 Background: Traditional Signal Propagation
When studying signal propagation in a network, the effect of a neural network layer on characteristics of the signal, such as its mean or variance, are investigated. The analysis of the vanishing gradient problem (Hochreiter, 1991) can also be considered as signal propagation theory where the norm of the delta-errors is tracked throughout layers in backpropagation. Typically, the network is assumed to be a repetitive composition of the same layers (e.g. fully-connected or convolutional). This allows to reduce the analysis of an entire network to a single layer and enables fixed-point analyses (e.g. Klambauer et al., 2017; Schoenholz et al., 2017). In our work, we focus on fully-connected networks,  $f: \mathbb{R}^N \to \mathbb{R}^M$ , with some activation function  $\phi: \mathbb{R} \to \mathbb{R}$  that is applied element-wise to vectors. We study the propagation of pre-activations² (cf. Saxe et al., 2014; He et al., 2015; Poole et al., 2016) throughout the network:
$$
\boldsymbol {s} = \boldsymbol {W} \phi \left(\boldsymbol {s} ^ {-}\right) + \boldsymbol {b}. \tag {1}
$$
Here,  $\mathbf{W} \in \mathbb{R}^{M \times N}$  and  $\mathbf{b} \in \mathbb{R}^{M}$  are the weight matrix and bias vector, respectively. The notation  $s^{-}$  is used to indicate the pre-activations from the preceding layer, such that  $s^{-} \in \mathbb{R}^{N}$ .
The first two moments of the signal can be expressed in terms of the randomness arising from the parameters. At initialisation time, the weight parameters are considered identically and independently distributed (i.i.d.) random variables  $w_{ij} \sim \mathcal{D}_w$ . The weight distribution is often assumed to be uniform or Gaussian, but the exact shape does not matter in wide networks (Golikov and Yang, 2022, Principle 2). The bias parameters, on the other hand, are commonly ignored in the analysis, which is justified when  $\forall i : b_i = 0$ . The pre-activations are also random variables due to the randomness of the parameters. If we assume the weights to be centred, such that  $\mathbb{E}[w_{ij}] = 0$ , and  $\mathrm{Var}[w_{ij}] = \sigma_w^2$ , the first two moments of the pre-activations are given by
$$
\mathbb {E} \left[ s _ {i} \right] = \mathbb {E} \left[ \sum_ {k} w _ {i k} \phi \left(s _ {k} ^ {-}\right) + b _ {i} \right] = 0 \tag {2}
$$
$$
\mathbb {E} \left[ s _ {i} ^ {2} \right] = \mathbb {E} \left[ \left(\sum_ {k} w _ {i k} \phi \left(s _ {k} ^ {-}\right) + b _ {i}\right) ^ {2} \right] = N \sigma_ {w} ^ {2} \mathbb {E} \left[ \phi \left(s _ {1} ^ {-}\right) ^ {2} \right]. \tag {3}
$$
Note that the variance is independent of the index,  $i$ , in the pre-activation vector and thus  $\forall k:\mathbb{E}\big[\phi (s_k^-)^2\big] = \mathbb{E}\big[\phi (s_1^-)^2\big]$ . Moreover, it can be shown (see Appendix A) that  $\forall i\neq j:\mathbb{E}[s_is_j] = 0$  i.e. features within a pre-activation vector are uncorrelated in expectation.
Signal propagation theory has been used to derive initialisation and normalisation methods (LeCun et al., 1998; Glorot and Bengio, 2010; He et al., 2015; Klambauer et al., 2017). Initialisation methods often aim at having pre-activations with the same mean and variance in every layer of the network. Because the mean is expected to be zero in every layer, the focus lies on keeping the variance of the pre-activations constant throughout the network, i.e.  $\forall i,j:\mathbb{E}\left[s_i^2\right] = \mathbb{E}\left[s_j^{-2}\right] = \sigma_*^2$ . Plugging this into Eq. (3), we obtain the fixed-point equation
$$
\sigma_ {*} ^ {2} = N \sigma_ {w} ^ {2} \operatorname {v a r p r o p} _ {\phi} \left(\sigma_ {*} ^ {2}\right). \tag {4}
$$
Here,  $\text{varprop}_{\phi}$  is a function that models the propagation of variance through the activation function,  $\phi$ , assuming zero mean inputs. E.g.  $\text{varprop}_{\text{ReLU}}(\sigma^2) = \frac{1}{2}\sigma^2$  (He et al., 2015) or  $\text{varprop}_{\text{tanh}}(\sigma^2) \approx \sigma^2$  (LeCun et al., 1998; Glorot and Bengio, 2010). We refer to Appendix A.3 for a more detailed discussion on propagation through activation functions.
This shows how modern initialisation methods are a solution to a fixed-point equation of the variance propagation. Our goal is to apply the same principles to find an initialisation strategy for the non-negative weights in ICNNs. However, Eq. (4) heavily relies on the assumption that the weights are centred, i.e.  $\mu_w = 0$ . This assumption is impossible to satisfy for the non-negative weights in ICNNs, unless all weights are set to zero. If  $\mu_w \neq 0$ , the mean of the pre-activations is no longer zero by default (as in Eq. 2). Furthermore, also the propagation of variance and covariance is affected. Therefore, we extend the signal propagation theory to accurately describe the effects due to non-centred weights.
${}^{2}$  a common alternative is to consider the (post-)activations (see Hoedt et al., 2018)
2.2 Generalised Signal Propagation
We generalise the traditional signal propagation theory by lifting a few assumptions from the traditional approach. Most notably, the weights are no longer drawn from a zero-mean distribution, such that  $\mathbb{E}[w_{ij}] = \mu_w \neq 0$ . Additionally, we include the effects due to bias parameters, which will give us extra options to control the signal propagation. Similar to the weights, we assume bias parameters to be i.i.d. samples from some distribution with mean  $\mathbb{E}[b_i] = \mu_b$  and variance  $\mathrm{Var}[b_i] = \sigma_b^2$ . By re-evaluating the expectations in Eq. (2) and Eq.(3), we arrive at the following results:
$$
\mathbb {E} \left[ s _ {i} \right] = N \mu_ {w} \mathbb {E} \left[ \phi \left(s _ {1} ^ {-}\right) \right] + \mu_ {b} \tag {5}
$$
$$
\mathbb {E} \left[ s _ {i} s _ {j} \right] = \delta_ {i j} \left(N \sigma_ {w} ^ {2} \mathbb {E} \left[ \phi \left(s _ {1} ^ {-}\right) ^ {2} \right] + \sigma_ {b} ^ {2}\right) + \mu_ {w} ^ {2} \sum_ {k, k ^ {\prime}} \operatorname {C o v} \left[ \phi \left(s _ {k} ^ {-}\right), \phi \left(s _ {k ^ {\prime}} ^ {-}\right) \right] + \mathbb {E} \left[ s _ {i} \right] \mathbb {E} \left[ s _ {j} \right], \tag {6}
$$
where  $\delta_{ij}$  is the Kronecker delta. We refer to Appendix A.2 for a full derivation. Note that mean and variance are still independent of the index,  $i$ . Therefore, we can continue to assume  $\forall k: \mathbb{E}\big[\phi (s_k^-)^n\big] = \mathbb{E}\big[\phi (s_1^-)^n\big]$ .
There are a few crucial differences between the traditional approach and our generalisation. First, the expected value of the pre-activations is no longer zero by default. This also means that the second moment,  $\mathbb{E}\big[s_i^2\big]$ , does not directly model the variance of the pre-activations in our setting. Secondly, the propagation of the second moment has an additional term that incorporates effects due to the covariance structure of pre-activations in the previous layer. Finally, the off-diagonal elements of the feature covariance matrix,  $\mathrm{Cov}[s_i,s_j] = \mathbb{E}_{i\neq j}[s_is_j] - \mathbb{E}[s_i]\mathbb{E}[s_j]$ , are not necessarily zero and also depend on the variance. This interaction between on- and off-diagonal elements in the feature covariance matrix makes it harder to ensure stable propagation.
Similar to Eq. (3), Eq. (6) depends on the variance propagation through the activation function. In addition, our generalised propagation theory involves the covariance propagation through the activation function. The covariance propagation through the ReLU function has been studied in prior work (Cho and Saul, 2009; Daniely et al., 2016) and can be specified as follows:
$$
\mathbb {E} \left[ \operatorname {R e L U} \left(s _ {1}\right) \operatorname {R e L U} \left(s _ {2}\right) \right] = \frac {\sigma^ {2}}{2 \pi} \left(\sqrt {1 - \rho^ {2}} + \rho \operatorname {a r c c o s} (- \rho)\right), \tag {7}
$$
with  $(s_1,s_2)\sim \mathcal{N}\bigg((0,0),\binom {1}{\rho 1}\sigma^2\bigg)$  and correlation  $\rho$ . Note that this expression implicitly describes the propagation of the squared mean  $(\rho = 0)$  and the second raw moment  $(\rho = 1)$ . We will assume  $\phi = \mathrm{ReLU}$  for the remainder of our analysis. A derivation for the propagation through the leaky ReLU activation function (Maas et al., 2013) can be found in Appendix A.3.
With our generalised propagation theory, the effects due to  $\mu_w \neq 0$  become apparent. As expected, the pre-activations no longer have zero mean by default. Furthermore, the covariance plays a significant role in the signal propagation. This shows that we also have to stabilise mean and covariance on top of the variance propagation to derive our principled initialisation.
3 Principled Initialisation of ICNNs
With our generalised signal propagation theory, we will now derive a principled initialisation strategy for ICNNs. First, we set up fixed point equations for mean, variance and correlation using our generalised framework. By solving these fixed point equations for  $\mu_w, \sigma_w^2, \mu_b$  and  $\sigma_b^2$  we obtain our principled initialisation.
3.1 Background: Input-Convex Neural Networks
Neural networks for which all input-output mappings are convex, are called Input-Convex Neural Networks (ICNNs) (Amos et al., 2017). In general, neural networks are functions that are constructed by composing one or more simpler layers or components. A function composition  $f \circ g$  is convex if both  $f$  and  $g$  are convex, where  $f$  additionally has to be non-decreasing in every output. Therefore, ICNNs are created by enforcing that every layer is convex and non-decreasing (Amos et al., 2017). A direct consequence is that only convex, non-decreasing activation functions can be
used in ICNNs — e.g. softmax (Dugas et al., 2001), ReLU (Nair and Hinton, 2010), LReLU (Maas et al., 2013), ELU (Clevert et al., 2016).
Fully-connected and convolutional layers are affine transformations, which are trivially convex. To make these layers non-decreasing for building ICNNs, their weights have to be non-negative (Amos et al., 2017). This can be done by replacing negative values with zero after every update. The non-negativity constraint does not apply to bias parameters. Another notable exception are direct connections from the input layer (Amos et al., 2017). This means that the first layer and any skip-connections from the input can have unconstrained weights. Amos et al. (2017) argue that these skip-connections are necessary for representational power. We limit our analysis to networks without skip-connections. Our experiments (see Section 5) show that ICNNs with our proposed initialisation do not require skip-connections for efficient training.
3.2 (Co-)Variance Fixed Points
To set up the fixed-point equation for the second moments (Eq. 6), we need to be able to propagate through the ReLU non-linearity. However, the analytical results from Eq. (7) only hold if the pre-activations have zero mean. Therefore, we look for a configuration that makes the mean (Eq. 5) of the pre-activations zero. Because we cannot enforce  $\mu_w = 0$ , we make use of the bias parameters to obtain centred pre-activations:
$$
\mu_ {b} = - N \mu_ {w} \mathbb {E} \left[ \phi \left(s _ {1} ^ {-}\right) \right]. \tag {8}
$$
This also allows us to refer to the second moments as the (co-)variance. Note that we approximate the pre-activations with Gaussians to use Eq. (7) at this point and for the remainder of the analysis. Figure 1c suggests that this Gaussian assumption generally does not hold in ICNNs (see also Appendix A.3). Nevertheless, our experiments (see Section 5) suggest that the approximation is sufficiently good to derive an improved initialisation strategy for ICNNs. Furthermore, Figure 6b indicates that the use of random initial bias parameters (cf. Appendix C.3) causes pre-activations to be (slightly) more Gaussian.
The propagation of the second moment described by Eq. (6) consists of two parts. The first part describes the off-diagonal entries, for which the dynamics are given by
$$
\mathrm {C o v} _ {i \neq j} [ s _ {i}, s _ {j} ] = \mu_ {w} ^ {2} \sum_ {k, k ^ {\prime}} \mathrm {C o v} [ \phi (s _ {k} ^ {-}), \phi (s _ {k ^ {\prime}} ^ {-}) ] = N \mu_ {w} ^ {2} \Big (\mathrm {V a r} [ \phi (s _ {1} ^ {-}) ] + (N - 1) \mathrm {C o v} [ \phi (s _ {1} ^ {-}), \phi (s _ {2} ^ {-}) ] \Big).
$$
Here, we rewrite the sum assuming that the pre-activations are identically distributed. This is possible because the covariance does not depend on the indices  $i$  or  $j$ . For further details, we refer to Appendix A.2 The second part models the on-diagonal entries, i.e. the variance, which can be simplified in a similar way
$$
\mathrm {V a r} [ s _ {i} ] = N \sigma_ {w} ^ {2} \mathbb {E} [ \phi (s _ {1} ^ {-}) ] + \sigma_ {b} ^ {2} + \mathrm {C o v} [ s _ {1}, s _ {2} ].
$$
By plugging the ReLU moments from Eq. (7) into these results in Appendix B.2, we obtain the desired fixed-point equations in terms of variance and correlation:
$$
\rho_ {*} = \mu_ {w} ^ {2} \frac {1}{2 \pi} N \left(\pi - N + (N - 1) \left(\sqrt {1 - \rho_ {*} ^ {2}} + \rho_ {*} \operatorname {a r c c o s} (- \rho_ {*})\right)\right) \tag {9}
$$
$$
\sigma_ {*} ^ {2} = N \sigma_ {w} ^ {2} \frac {1}{2} \sigma_ {*} ^ {2} + \sigma_ {b} ^ {2} + \rho_ {*} \sigma_ {*} ^ {2}. \tag {10}
$$
with  $\rho_{*} = \frac{1}{\sigma_{*}^{2}}\mathrm{Cov}[s_{1}^{-},s_{2}^{-}] = \frac{1}{\sigma_{*}^{2}}\mathrm{Cov}_{i\neq j}[s_{i},s_{j}]$  and  $\sigma_{*}^{2} = \mathrm{Var}[s_{1}^{-}] = \mathrm{Var}[s_{i}]$ .
3.3 Weight Distribution for ICNNs
To obtain the distribution parameters for the initial weights, we solve the fixed point equations in Eq. (9) and Eq.(10) for  $\sigma_b^2$ ,  $\sigma_w^2$  and  $\mu_w^2$ . Because this system is over-parameterised, we choose to set  $\sigma_b^2 = 0$ . An analysis for  $\sigma_b^2 > 0$  can be found in Appendix C.3. This leads to the following initialisation parameters for ICNNs:
$$
\sigma_ {w} ^ {2} = \frac {2}{N} (1 - \rho_ {*}) \tag {11}
$$
$$
\mu_ {w} ^ {2} = \frac {2 \pi}{N} \rho_ {*} \left(\pi - N + (N - 1) \left(\sqrt {1 - \rho_ {*} ^ {2}} + \rho_ {*} \operatorname {a r c c o s} (- \rho_ {*})\right)\right) ^ {- 1}. \tag {12}
$$
Note that these solutions only depend on the correlation,  $\rho_{*}$ , and not on the variance,  $\sigma_{*}^{2}$ . For a stability analysis of these fixed points, we refer to Appendix B.4.
Although our initialisation (Eq. 12 and 11) is entirely defined by the correlation between features,  $\rho = \frac{\mathrm{Cov}[s_1,s_2]}{\mathrm{Var}[s_1]}$ , not all solutions are admissible. For example, uncorrelated features  $(\rho = 0)$  are only possible if  $\mu_w = 0$ . This means that features in ICNNs must have non-zero correlation. On the other hand, perfect correlation  $(\rho = 1)$  would mean that all features point in the same direction, which is not desirable. Therefore, we choose  $\rho_{*} = \frac{1}{2}$  as a compromise to obtain the following initialisation parameters for our experiments in Section 5:
$$
\mu_ {w} = \sqrt {\frac {6 \pi}{N (6 (\pi - 1) + (N - 1) (3 \sqrt {3} + 2 \pi - 6))}} \quad \sigma_ {w} ^ {2} = \frac {1}{N}
$$
$$
\mu_ {b} = \sqrt {\frac {3 N}{6 (\pi - 1) + (N - 1) (3 \sqrt {3} + 2 \pi - 6)}} \qquad \sigma_ {b} ^ {2} = 0.
$$
We refer to appendix C.4 for a discussion on different choices for  $\rho_{*}$ .
For the first layer, which can also have negative weights, we use LeCun et al. (1998) initialisation. For the non-negative weights in an ICNNs, we need to sample from a distribution with non-negative support. The lower bound of a uniform distribution with our suggested mean and variance is negative for any  $N > 1$  and thus not useful. Also, sampling from a Gaussian distribution is not practical because its support includes negative values. Eventually, we propose to sample from a log-normal distribution with parameters
$$
\tilde {\mu_ {w}} = \ln (\mu_ {w} ^ {2}) - \frac {1}{2} \ln (\sigma_ {w} ^ {2} + \mu_ {w} ^ {2}) \quad \quad \quad \quad \tilde {\sigma_ {w}} ^ {2} = \ln (\sigma_ {w} ^ {2} + \mu_ {w} ^ {2}) - \ln (\mu_ {w} ^ {2}).
$$
This ensures that sampled weights are non-negative and have the desired mean and variance. Note that other positive distributions with sufficient degrees of freedom should also work.
One advantage of the log-normal distribution is that sampling is simple and efficient. It suffices to sample  $\tilde{w}_{ij} \sim \mathcal{N}(\tilde{\mu_w}, \tilde{\sigma_w}^2)$  to compute  $w_{ij} = \exp(\tilde{w}_{ij})$ . In the original ICNN (Amos et al., 2017), the exponential function is only used to make the initial weights positive. However, if weights are re-parameterised using the exponential function, initial weights can be directly sampled from a Gaussian distribution. This setting is studied in Appendix C.5.