Principled Weight Initialisation for Input-Convex Neural Networks
Authors: Pieter-Jan Hoedt
Günter Klambauer
Venue: NeurIPS 2023 poster (2023)
Keywords: initialization
signal propagation
input-convex networks
TL;DR: We generalise signal propagation theory to derive an initialisation for the positive weight matrices in ICNNs

Abstract: 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.

Principled Weight Initialisation for Input-Convex Neural Networks
Pieter-Jan Hoedt & Günter Klambauer
LIT AI Lab & ELLIS Unit Linz Institute for Machine Learning Johannes Kepler University, Linz, Austria {hoedt, klambauer) j.u. at
Abstract
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 achieved by using non-decreasing convex activation functions and non-negative weights. Because of these peculiarities, previous ini- tialisation strategies, which implicitly assume centred weights, not effective for ICNNs. By studying signal propagation through layers with non-negative weights, we are able derive principled weight initialisation for ICNNs. Con- cretely, we generalise signal propagation theory by removing the assumption that weights are sampled from a centred distribution. In 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 common belief, ICNNs can be trained without skip-connections when initialised correctly. Finally, we apply ICNNs to real-world drug discovery task and show that they allow for more effective molecular latent space exploration.
1 Introduction
Input-Convex Networks. Input-Convex Neural Networks (ICNNs) are networks for which each output neuron is convex with respect to the inputs. The convexity 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 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.. inputs for which the output prediction re- mains unchanged. ICNNs have also been used in model predictive control, where recently 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. (d)). We hypothesize that this slow training is the result of poor initialisation of the positive weights in ICNNs. This poor initialisation leads to distri- bution 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 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.,
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

pre-activation distributions
055
0.5
42
weight signs or 2
feature correlations
(a) traditional
(b) ICNN
(c) ICNN init
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.
2020) or during back-propagation (Glorot and Bengio, 2010; Hoedt et al., 2018 Defazio and Bottou, 2021). It is also important account for the effects due to non-linearities in the network (Saxe al.,
2014; He et al., 2015; Klambauer et al., 2017: Hoedt et al., 2018) 80. 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 prop- agation 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 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 be- cause some weight matrices constrained be non-negative (Amos al., 2017). Therefore, 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 increase their representational power (Amos al., 2017). Although possible to study signal propagation with skip-connections (e.g. Yang and Schoen- holz, 2017; Brock et al., 2021; Hoedt al., 2022), they typically built on top of existing results for plain networks. Therefore, we limit our theoretical results to ICNNs without skip-connections. However, find that ICNNs without skip-connections skip-connections might not be necessary for representational power. We show that with our initiali- sation, 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).
be successfully trained, indicating that
Contributions. Our contributions can be summarised as follows:
We generalise signal propagation theory to include weights without zero mean (Section 2).
We derive principled initialisation strategy for ICNNs from our new theory (Section 3).
We empirically demonstrate the effectiveness of our initialisation strategy (Section 5.
We apply ICNNs in real-world -iscodeyy 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 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 derive a weight initialisation strategy for ICNNs (Section 3).
Code for figures and experiments can be found https: www.ccommmmm..ee..c wcwm/mll-jk//ccooe---inn

---

2.1 Background: Traditional Signal Propagation
When studying signal propagation in a network, the effect of a neural network layer on character- istics of the signal, such as its mean or variance, 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 single layer and enables fixed-point anal- yses (e.g. Klambauer et al., 2017; Schoenholz et al., 2017). In our work, we focus on fully-connected networks, f EV + R with some activation function & R + R that is applied element-wise to vectors. We study the propagation of pre "reaactivations" (cf. Saxe et al., 2014; He et al., 2015; Poole al., 2016) throughout the network:
investigated. The analysis of the vanishing
8 ( +b
(1)
Here, W MxN and b RM are the weight matrix and bias vector, respectively, The notation 8 is used to indicate the pre-activations from the preceding layer, such that & E RR.
The first two moments of the signal can be expressed in terms of the randomness arising from the pa- rameters. At initialisation time, the weight parameters are considered identically and independently distributed (i.i.d.) random variables Wi; Dw. The weight distribution is often assumed to be uni- form 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 Vi:: : by = 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 E/way] 0, and Var (may o the first two moments of the pre-activations are given by
E [ss:]
E
wan 6 bb
+b6
w + bi) bi) = No [(o(s)))))))
E
E
(3)
Note that the variance is independent of the index, i, in the pre-activation vector and thus Vk- E [[((5 } = E [( Moreover, it can be shown (see Appendix A) that Vi + j E[s;s,] = 0,
ie. 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. Vi,j E [+]] into Eq. (3), we obtain the fixed-point equation
E [s] = a Plugging this
0² No varprop, (o²).
(4)
Here, varprops is a function that models the propagation of variance through the activation function, o, assuming zero mean inputs. E.g. varproPraLu(o") varproptant (o²) x x² (LeCun al., 1998; Glorot and Bengio, 2010). We refer to Appendix A.3 for a more detailed discussion on propagation through activation functions.
_ (He et al., 2015) or
This shows how modern initialisation methods are solution to fixed-point equation of the variance propagation. Our goal 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. Pw
impossible to satisfy for the non-negative weights in ICNNs, unless all weights are set to zero. If ww A 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.
This assumption
a common alternative is to consider the (post- Jactivations (see Hoedt et al., 2018)
3

---

2.2 Generalised Signal Propagation
We generalise the traditional signal propagation theory by lifting few assumptions from the tradi- tional approach. Most notably, the weights are longer drawn from zero-mean distribution, such that E[w,y] = fww # O. Additionally, we include the effects due to bias parameters, which will give us extra options control the signal propagation. Similar to the weights, to i..d. samples from some distribution with mean E[b,] μs and variance Var[b:] oi. By re-evaluating the expectations in Eq. (2) and Eq.(3), we arrive at the following results:
assume bias parameters
E[s:] N Mm ( 1AB
(5)
2 E [ss d.y (No E [ + pr
Σ Corlo(s7). [[ss] E[sy], k.k'
(6)
where di 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, can continue to assume Ak E [6($)))) = =B [()))))
There are 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, [ does not directly model the variance of the pre-activations in 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, Cov [8,, 85 [* [E[s,] R[sy). 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 Eq. (3), Eq. (6) depends 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:
0² (V1 - p² parccos(-p)).
ReLLU(s)))
2
N (0.0), (b and correlation Note that this expression implicitly de-
with (s1, No
scribes the propagation of the squared mean (p = 0) and the second raw moment (p = 1). We will assume Φ ReLU for the remainder of 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 Muw * 0 become apparent. As expected, the pre-activations no longer have zero mean by default. Furthermore, the covariance plays signif- icant role in the signal propagation. This shows that 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 principled initialisation strat- egy 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 pows o2 H4, and of 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 con- structed by composing one or more simpler layers or components. function composition I is convex if both f and are convex, where 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 eeg.. softplus (Dugas et al., 2001), ReLU (Nair and Hinton, 2010), LReLU (Maas al., 2013), ELU (Clevert 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 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 bias parameters. Another notable exception are direct connections from the imput layer (Amos et al., 2017). This means that the first layer and any skip- connections from the input can have unconstrained weights. Amos al. (2017) argue that these skip ssip-connections are necessary for representational power. We limit our analysis to networks without skip ciipeconnections. Our experiments (see Section show that ICNNs with our proposed initialisation do not require coppeenneetinns for efficient training.
3.2 (Co-anariance Fixed Points
To up the fixed-point equation for the second moments (Eq. 6), need to be able propagate through the ReLU non-linearity. However, the analytical results from Eq. (7) only hold the pre- activations have zero mean. Therefore, we look for configuration that makes the mean (Eq. 5) of the pre-activations zero. Because we cannot enforce Va O, we make use of the bias parameters to obtain centred pre-activations:
Hb = -N F[
(8)
This also allows refer to the second moments the (co-)variance. Note that approximate the pre-activations with Gaussians to use Eq. (7) at this point and for the remainder of the analy- sis. Figure Ic 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 derive 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
Cove s5] = Niw Σ (oo(((s)).
Nr (Vrlo(ss ) + +NV 1) Covlo(s)).
Here, we rewrite the sum assuming that the pre-activations are identically distributed. This is pos- sible 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 a similar way
Var[s:] = No B[o(sT)] )] + of + Cov[s1,82]
By plugging the ReLU moments from Eq. (7) into these results Appendix B.2, we obtain the desired fixed-point equations in terms of variance and correlation:
P. = M 1w
N
N (N ()(1 pp2
22 + arccos(-p.))
(9)
2m
- = No 3
o. - + of + p.o
(10)
with p. o. Cov(sT s5] +/2 Covizylsis sj] and of = Var[s]] Var[s:].
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 (10 for of o² and a Because this system is over-parameterised, we choose to set of = O. An analysis for of > can be found in Appendix C.3. This leads to the following initialisation parameters for ICNNs:
2
a (1 -p.) P.)
(11)
N
2n
P N
P. -
N NN 1)(11 p2 (rccos(-pp)))
(12)
5

---

Note that these solutions only depend on the correlation, Px. and not on the variance, o. For 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, Varlei not all solutions are admissible. For example, uncorrelated features (p 0) are only _sol
possible if Pru O. This means that features in ICNNs must have non-zero correlation. On the other hand, perfect correlation (p= 1) would mean that all features point in the same direction, which is not desirable. Therefore, choose P. 2 parameters for our experiments in Section 5:
compromise to obtain the following initialisation
6T
1 t N
(6(T 11) (N- 1)(3/3 + 2m 6))
3NV
14 V 6(TT 1) (N - 1)3333 2T -6)
o = 0.
We refer to appendix C.4 for a discussion on different choices for Pa-
For the first layer, which 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 sample from log-normal distribution with parameters
Fu = In( In + +
F = + + - In(
This ensures that sampled weights are non-negan 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 lor-normal distribution is that sampling simple and efficient. It suffices to sample úij ~ (( ow to compute Way exp( (.y). 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 Gaussian distribution. This setting studied in Appendix C.5.
4 Related Work
Input-conver neural networks. Input-Convex Neural Network (ICNN) were originally designed in the context energy models (Amos et al., 2017). We focus on the fully convex ICNNs variant and use gradient descent for optimisation instead of the proposed bundle-entropy method. In their implementation, Amos et al. (2017) use projection methods to keep weights positive, i.c. negative values are set to zero after every update. However, also other projection methods can be used to keep the weights positive (e.g. Sivaprasad et al., 2021). Instead of projecting the weights onto the non-negative reals after every update, it is also possible use reparameterisation of the weights. E.g. Nesterov et al. (2022) square the weights in the forward pass. Note that a reparameterisation can have an effect on the learning dynamics because it is inherent part of the forward pass. An- other alternative is to use regularisation to impose a soft non-negativity constraint on the weights (Makkuva et al., 2020). Similar (Sivaprasad al., 2021), we directly train ICNN, but do not observe the same generalisation benefits. On the other hand, Sankaranarayanan and Ren- gaswamy (2022) point out that restricting the network to be convex decreases their capacity. They propose to use the difference of two ICNNs to allow modelling non-convex functions (c.f. Yuille and Rangarajan, 2001). However, we did not find the convexity restriction to cause any problems in our experiments. Nesterov et al. (2022) use ICNNs to simplify the computation of level sets to learn invariances. We adopt this experiment setting to highlight the potential of ICNNs. Optimization. While training the parameters of an ICNN is unrelated to convex optimisation, Bengio et al. (2005) showed how training a neural network together with the number of hidden neurons can be seen as a convex optimisation problem. This work was continued by Bach (2017), who established connec- tion between this convex optimisation and automatic feature selection. Signal propagation theory. The study of signal propagation is concerned with tracking statistics of the data throughout multiple

---

layers of the network. In his thesis, Neal (1995) computed how the first two moments propagate through a two-layer network to study networks as Gaussian processes. This is similar to how mod- em initialisation strategies have been derived (e.g. LeCun et al., 1998; He et al., 2015). A similar approach has been taken to incorporate the propagation of gradients for initialisation (Glorot and Bengio, 2010). The idea to explicitly account for the effects of activation functions can be attributed to (Saxe et al., 2014). All of these analyses assume that weights are initialised from a zero-mean distribution, which is not possible in ICNNs. Mishkin and Matas (2016) empirically compute the variance for the weights. This approach is more robust to deviations, but nevertheless requires adap- tations to account for weights with non-zero mean. There is also significant body of work that studies signal propagation in terms of mean field theory (Poole et al., 2016). In these works, the correlation between samples is included in the analysis (Poole et al., 2016). In our work, find that the correlation between features plays an important role. Schoenholz et al. (2017) applied the mean field analysis to the backward dynamics and introduced the concept of depth scales to quantify how deep signals can propagate. The mean field theory was also applied to convolutional networks to derive the delta-orthogonal initialisation (Xiao et al., 2018). We refer to (Martens et al., 2021) for a general overview of signal propagation.
5 Experiments
We evaluate our principled initialisation strategy by training ICNNs on three sets of experiments. In our first experiments, we investigate the effect of initialisation on learning dynamics and gener- alisation in ICNNS on multiple permuted image datasets. We also include non-convex networks in these experiments illustrate that ICNNs with our principled initialisation can be trained as well as regular networks. However, we would like to stress that non-conve networks are expected to outper- form ICNNs because they are not constrained to convex decision boundaries. Moreover, we assess whether ICNNs using our initialisation are able to match the performance of regular fully-connected networks in toxicity prediction without the need for skip-connections. Finally, we explore ICNNS as tool for more controlled latent space exploration of a molecule auto-encoder system. This illus- trates a setting where regular networks can no longer be used. Details on the computing budget are provided in the Appendix.
5.1 Computer vision benchmarking datasets
As an illustrative first set of experiments, we use computer vision benchmarks to assess the effects of our principled initialisation strategy. Concretely, we trained fully-connected ICNNs on MNIST (Bottou et al., 1 1994), CIFARIO and CIFAR (Krizhevsky, 2009), to which we refer as permuted image benchmarks (ef. Goodfellow et al., 2014).
In our comparison, we consider four different types of methods: (1) a classical non-convex network without skip ikip-connections and default initialisation (ef. He et al., 2015). (2) an ICNN, "ICNN", with- out skip connections and default initialisation (3) an ICNN, "ICNN + skip", with skip connections and default initialisation (4) an ICNN, "ICNN init", without skip connections and our proposed initialisation, All networks use ReLU activation functions. The non-negativity constraint in ICNNS implemented by setting negative weights zero after every update (cf. Amos et al., 2017). Note that skip- conpeconnestions in ICNNs introduce additional parameters to allow for additional representa- tional power (cf. Amos et al., 2017).
Training dynamics. Setings. In the first experiment, we analyze the training dynamics during the first ten epochs of training. Similar to Chang et al. (2020), we fixed the number of neurons in each layer to the input size, and the number of hidden layers to five. We refer to Appendix C.2 for results with different depths. The learning rate for the Adam optimiser on the non-convex baseline. Results. We found that ICNNS with our principled initialisation ex- hibit similar learning behaviour the non-convex baseline, while ICNNs without our initialisation strategy can not decrease the loss to the level of the baseline on CIFARIO and DEFAR100, Figure shows the learning curves for our different methods on three permuted image benchmarks.
obtained by manually tuning
Generalisation. Settings. In our next experiment, the effects of our principled initialisation on the generalisation capabilities of ICNNs is investigated. To this end, we train our three ICNN variants and the non-convex baseline on the permuted image benchmarks again, but focus on the test perfor- that attain the
mance this time. To this end we perform a grid-search to find the
7

---

CIAR30
10
127
1250
107
337 2200
2030
0000 20.00
100
400
1000
Figure 2: Training loss curves of ICNN variants with the same architecture on the MNIST. CIFARIO and CIFARIOO 100 datasets. "ICNN" input-convex network with default initialisation. "ICNN + skip": same settings, but with skip connections. "ICNN + init" our principled initialisation for ICNNs w/o skip connections. "non-convex": a regular non-convex network. Each curve represents the median performance over ten runs and shaded regions indicate the inter-quartile range. The inset figures provide view of the loss curves zoomed in. Note that ICNN losses do decrease before the plateau.
15'
150
1-11
ICH
Elat- /M
039
6.0
26
14 20 25
5
15
24
Figure 3: Test set metrics of compared methods on the (a) MNIST and (b) CIFARIO datasets. Each curve displays the median performance over ten runs Shaded regions represent the inter-quartile range over the runs. On MNIST, all methods be successfully trained and exhibit similar per- formance. On CIFAR 10, ICNNs with our proposed initialisation outperform other ICNNs variants.
best accuracy after 25 epochs of training on a random validation split, for each of the four compared methods. The grid consists of multiple architectures with at least one hidden layer, learning rate for Adam, Le-reullarisation and whitening pre-processing transforms (details in Appendix C.2). Re- sults. On the MNIST dataset, all variants reach similar accuracy values, which we attribute to the simplicity of the prediction problem. However, on CIFAR 10 the initialisation strategy leads to better generalisation already in early epochs and compared to the respective ICNNs variants without our initialisation improves generalisation overall (see Figure 3). The ICNN with our initialisation almost matches the accuracy values of non-convex nets.
5.2 Toxicity Prediction
To test ICNNs in a real-world setting, in a different application domain, we consider the binary multi-task problem of toxicity prediction in drug discovery. More specifically, train ICNNs on the Tox21 challenge data (Huang et al., 2016; Mayr et al., 2016; Klambauer et al., 2017). The input data consists of so-called SMILES representations of small molecules. The target variables are binary labels for twelve different measurements or assays indicating whether a molecule induces a particular toxic effect on human cells. During pre-processing, SMILES strings were converted to Continuous and Data-Driven Descriptors (CDDDs) (Winter et al., 2019), which are 5 512-dimensional numerical representations of the molecules. The choice of these particular descriptors is motivated by CDDDs properties of being decodeable into valid molecules (see Section 5.3). Hyper-parameters were selected by manual search on the non-convex baseline. We ended up using a fully-connected network with two hidden layers of 128 neurons and ReLU activations. The network was regularised with fifty percent dropout after every hidden layer, as well as seventy percent dropout of the inputs. The results in Table show that our initialisation significantly outperforms ICNNs with standard initialization (p-value 2.6e- 13, binomial test) and ICNNs with skip connections (p-value 5.5e-14). Furthermore, results are close to the performance of traditional, non-convex networks.
8

---

Table 1: Area under the ROC curve on the test set for each of the twelve tasks in the Tox21 data. Each value represents the median performance over 10 runs and eπror bars is the maximum distance to the boundary of the interval defined by the (0.05, 0.95) quantiles. Of variants of ICNNs, the variant with our proposed initialisation, "ICNN+init (ours)", " performs best matching almost the predictive quality of traditional non-convex neural networks.
NR.AR
LBD
NR.Aromatase
ICNNsinic (ouns)
2.40%
78.37
ICNN
90.56 0.63%
62.28 4: 4.52%
78.14
5.18%8
30.40 4: 1.74%
17.30
0.78%
78.25
1.002%
ICNN-skip
89.83 4: 0.21%
68.12 1.74%
71.17
4. 2.20%
78.95
0. 45%
70.95
0.48%
81.92 2: 16%
0.053
79.21
1.70%
2.50%
88.63
0.3456
27.88
1.13%
74.32
1.50%
SR.ATADS
SR.HSE
SR MMP
SR.p53
PPARY
AVG
IONN-init (ours)
77:01
1.67%
N0.05 + 1.36%
0 b.27%
0.46%
77.46
2.00%
0.33%
80.57
ICNN
74.25 + 3.26%
78.64 1.4256
93.19
0.50%
$7.92
0.4955
75.00
1.73%
76.11 0.82%
78.33
13.425
ICNNsskip
75.99
1.26%
78.23
0.50%
75.40
0.88/%
78.25 4: 0.93%
92.09 4: 0.41%
79.32 d: 0.97%
18.29
12.10%
1.68%
2.57%
0.30%
0.96%
0.90%
81.31
11.08%
6003
QRf 6043
(IT 0/00
NILSUR 2018
NIEELED 20T8
INTRREE DEIS
NESSUAR 2042
Figure 4: Example of two level set trajectories of a Tox21 model. The leftmost molecule represents the starting molecule with low toxicity (top: "NR.AR", bottom: "NR.ER.LBD"), The rightmost molecule represents the target molecule. All intermediary molecules are samples on the level set between the starting molecule and the target molecule and exhibit different drug-likeness ("QED"). In this way, one predicted molecular property can be kept fixed while another property is optimized.
5.3 Latent space exploration of molecular spaces
In this experiment, we exploit the intrinsic property of ICNNS that level sets, i.e. the set of all inputs that map to the same output, can be parameterised (Nesterov et al., 2022). These level sets provide an opportunity drug discovery to keep one molecular property, such low toxicity, fixed while optimizing another property, such as drug-likeness. To demonstrate this application, we randomly chose reference molecules with low toxicity and followed the according level set the direction of another randomly chosen target molecules. Since the input space of the ICNN that predicts toxicity (see above) possesses decoder model from (Winter et al., 2019), the numeric representations of the molecules on the level set be decoded into molccular structures. The obtained trajectories in Figure demonstrate that the predicted toxicity of the molecules remain constant, while different QED scores, which measure drug- likeness, are obtained by traversing the level set from the reference molecule to the target molecule.
6 Conclusion and Discussion
We demonstrated how signal propagation theory can be generalised to include weights without zero mean. By controlling the propagation of correlation through the network, we derived an initialisation strategy that promotes stable propagation of signals. We empirically verified that this steady prop- agation leads improved learning. Finally, showed that ICNNs can be trained to comparable level as regular networks on the Tox21 data using CDDDs.
Our initialisation approximates the distribution of pre-activations with Gaussians, which falls short of agreement with the observed data, but has been sufficient to improve the initialisation ICNNs. A possible way forward is to use the central limit theorem for weakly dependent variables. In Ap-

---

pendix C.3, we also observed that random bias initialisation leads to pre-activations with a more Gaussian distribution. Appendix A.5 also includes an analysis for convolutional layers, but profi- cient empirical results with input-convex convolutional networks are still to be made. We conjecture that this is due to the more complicated covariance structure in convolutional layers. We also in- cluded signal propagation analysis for back-prov proppaattn in Appendix A.4. However, incorporating insights from this analysis into an initialisation method is also left for future work. We envision that our initialisation strategy will make it easier to incorporate ICNNs or other networks with non- negative weights.
Acknowledgments and Disclosure of Funding
We thank Sepp Hochreiter for pointing us in the direction of input-convex networks and their in- teresting properties. Special thanks go to Niklas Schmidinger for his help with experiments on convolutional networks.
The ELLIS Unit Linz, the LIT AI Lab, the Institute for Machine Learning, are supported by the Federal State Upper Austria. We thank the projects AI-MOTION (LIT-2018-6- YOU-212), DeepFlood (1T-2201---YOU-222 INCONTROL-RL. (FFG-881064), PRIMAL (FFG-873979). S3AI (FFG-872172), DL for Gran- ularFlow (FFG-871302), EPILEPSIA (FFG-892171), AIRI FG 9-N (FWF-36284, FWF-36235), AMGreenfFeatingGidssTTTG 899943), INTEGRATE (FFG-892418), ELISE (H2020-ICT-2019- DARIZZON-CL-222--LLMAATTII
Medical Cognitive Computing Center (MC3),
3 ID: 951847), Stars4Waters
We thank AudiJJKU Deep Learning Center, TGW LOGISTICS GROUP GMBH, Silicon Austria Labs (SAL), FILL Gesellschaft mbH, Anyline GmbH, Google, ZF Friedrichshafen AG, Robert Bosch GmbH, UCB Biopharma SRL, Merck Healthcare KGaA, Verbund AG, GLS (Univ. Waterloo) Software Com- petence Center Hagenberg GmbH, TÚV Austria, Frauscher Sensonic, TRUMPF and the NVIDIA Corporation.
References
Amos, B., Xu, L., and Kolter, J. Z. (2017). Input convex neural networks. In Proceedings of the 34th International Conference on Machine Learning, volume 70, pages 146-155. PMLR.
Bach, F. (2017). Breaking the curse of dimensionality with convex neural networks Journal of Machine Learning Research, 18(19):1-53.
Bengio, Y., Roux, N., Vincent, P., Delalleau, O., and Marcotte, P. (2005). Convex neural networks. In Advances Neural Information Processing Systems, volume 18, pages 123-130. MIT Press.
Bottou, L., Cortes, C., Denker, J., Drucker, H., Guyon, I., Jackel, L., LeCun, Y., Muller, U., Sackinger, E., Simard, P., and Vapnik, V. (1994). Comparison of classifier methods: a case study handwritten digit recognition. In Proceedings of the 12th IAPR International Conference on Pattern Recognition, volume 2, pages 77-82, Jerusalem, Israel. IEEE.
Brock, A., De, S., and Smith, S. L. (2021). Characterizing signal propagation to close the perfor- mance gap in unnormalized ResNets. In International Conference on Learning Representations, volume 9, virtual. arXiv: 10.69.....
Chang, O., Flokas, L., and Lipson, H. (2020). Principled weight initialization for hypernetworks. In International Conference Learning Representations, volume 8.
Cho, Y. and Saul, L. (2009). Kernel methods for deep learning. Advances in neural information processing systems, 22.
Clevert, D.-A., Unterthiner. T., and Hochreiter, S. (2016). Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs). International Conference on Learning Representations, volume 4, San Juan, Puerto Rico. arXiv: 1511.07289.
Daniely, A., Frostig, R., and Singer, Y. (2016). Toward Deeper Understanding of Neural Networks: The Power of Initialization and Dual View on Expressivity. Advances in Neural Information Processing Systems, volume 29, pages 2253-2261 Barcelona, Spain. Curran Associates, Inc.
10

---

Defazio, A. and Bottou, L. (2021). Beyond Folklore: A Scaling Calculus for the Design and Initial- ization of ReLU Networks. arXiv: 1906.04267.
Dugas, C., Bengio, Y. Bélisle, F., Nadeau, C., and Garcia, R. (2001). Incorporating Second-Order Functional Knowledge for Better Option Pricing. In Advances in Neural Information Processing Systems, volume 13, pages 472-478, Vancouver, BC, Canada. MIT Press.
Glorot, X. and Bengio, Y. (2010). Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9, pages 249-256.
Golikov, E. and Yang, G. (2022), Non-gaussian tensor programs. In Oh, A. H., Agarwal, A., Bel- grave, D., and Cho, K., editors, Advances in Neural Information Processing Systems, volume 35.
Goodfellow, J., Mirza, M., Da, X., Courville, A. C., and Bengio, Y. (2014). An Empirical Inves- tigation of Catastrophic Forgeting in Gradient-Basse Bassed Neural Networks. In International Confer- ence on Learning Representations, volume 2.
He, K., Zhang, X., Ren, S., and Sun, J. (2015). Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026-1034.
Hochreiter, S. (1991). Untersuchungen zu dynamischen neuronalen Netzen. Master's thesis, Tech- nical University Munich, Munich, Germany.
Hoedt, P.-J., Hochreiter, S., and Klambauer, G. (2018). Characterising activation functions by their backward dynamics around forward fixed points. NeurlPS workshop on Critiquing and Correcting Trends Machine Learning.
Hoedt, P.-.., Hochreiter, S., and Klambauer, G. (2022). Normalization is dead, long live normaliza- rresssseemmmmiiirrrreeeessss.
tion! In ICLR Blog Track. https://iclr-blog wwloogotccc.giieeeee
Huang, R., Xia, M., Nguyen, D.-T., Zhao, T., Sakamuru, S., Zhao, J., Shahane, S. A., Rossoshek, A., and Simeonov, A. (2016). Tox21Challese to Build Predictive Models of Nuclear Receptor and Stress Response Pathways as Mediated by Exposure to Environmental Chemicals and Drugs Frontiers in Environmental Science,
Klambauer, G., Unterthiner, T., Mayr, A., and Hochreiter, S. (2017). Seff-normalizing neural net- works. Advances in neural information processing systems, 30.
Krizhevsky, A. (2009). Learning Multiple Layers of Features from Tiny Images. Technical report, Canadian Institute for Advanced Research.
LeCun, Y., Bottou, L., Orr, G. B., and Müller, K. R. (1998). Efficient BackProp. In Orr, G. B. and Müller, K.-R., editors, Neural Networks: Tricks of the Trade, Lecture Notes in Computer Science, pages 9-50. Springer, Berlin, Heidelberg, 1 edition.
Maas, L., Hannun, A. Y., and Ng. A. Y. (2013). Rectifier Nonlinearities Improve Neural Network Acoustic Models. In Workhop for Deep Learning for Audio, Speech and Language Processing at ICML, Atlanta, GA. USA.
Makkuva, A., Taghvaei, A., Oh, S., and Lee, J. (2020). Optimal transport mapping via input convex neural networks. In Proceedings of the 37th International Conference Machine Learning, volume 119, pages 6672-6681 6681. PMLR.
Martens, J., Ballard, A., Desjardins, G., Swirszez, G., Dalibard, V., Sohl-Dickstein, J., and Schoen- holz, S. S. (2021). Rapid training of deep neural networks without skip connections or normal- ization layers using Deep Kernel Shaping. arXiv: 21 10.01765.
Mayr, A., Klambauer, G., Unterthiner, T., and Hochreiter. S. (2016). Deeptox: toxicity prediction using deep learning. Frontiers Environmental Science, 3:80.
Mishkin, D. and Matas, J. (2016). All you need good init. In International Conference on Learning Representations, volume 4.
11

---

Nair, V. and Hinton, G. E. (2010). Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th Interaational Conference on Machine Learning, pages 807-814, Madi- Son, WI, USA. Omnipress.
Neal, R. M. (1995). Bayesian Learning for Neural Networks. PhD thesis, University of Toronto.
Nesterov, V., Torres, F. A., Nagy-Huber, M., Samarin, M., and Roth, V. (2022). Learning invariances with generalised input-convex neural networks.
Poole, B., Lahiri, S., Raghu, M., Sohl-Dickstein, J., and Ganguli, S. (2016). Exponential expressivity in deep neural networks through transient chaos. In Advances in Neural Information Processing Systems, volume 29, pages 3360-3368.
Sankaranarayanan, P. and Rengaswamy, R. (2022). CDiNN convex difference neural networks. Neurocomputiig, 495:153-168.
Saxe, A., McClelland, J., and Ganguli, (2014). Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In International Conference Learning Representa- tions, volume 2.
Schoenholz, S. S., Gilmer, J., Ganguli, S., and Sohl-Dickstein, J. (2017). Deep information propa- gation. In International Conference on Learning Representations.
Sivaprasad, S., Singh, A., Manwani, N., and Gandhi, V. (2021). The curious case of convex neural networks. In Oliver, N., Pérez-Cruz, F., Kramer, S., Read, J., and Lozano, J. A., editors, Machine Learning and Knowledge Discovery in Databases. Research Track, Lecture Notes in Computer Science, pages 738-754. Springer International Publishing.
Wang, Z. and Wu, Z. (2023). Input convex Istm: A convex approach for fast Lyapunov-based model predictive control. arXiv preprint arXiv:3311.07202.
Winter, R., Montanari, F., Noć, F., and Clevert, D.-A. (2019). Learning continuous and data- driven molecular descriptors by translating equivalent chemical representations, Chemical Sci- ence, 1666)1192--7711. Publisher: The Royal Society of Chemistry.
Xiao, L.., Bahri, Y., Sohl-Dickstein, J., Schoenholz, S., and Pennington, J. (2018). Dynamical isometry and mean field theory of cnns: How train 10,000-layer vanilla convolutional neural networks. In Proceedings of the 35th International Conference on Machine Learning, volume 80, pages 5393-5402. PMLR.
Yang, G. and Schoenholz, S. (2017). Mean field residual networks: On the edge of chaos. Advances in neural information processing systems, 30.
Yuille, A. L. and Rangarajan, A. (2001). The concave-convex procedure (cccp). In Advances in Neural Information Processing Systems, volume 13, pages 1033-1040.
Zhang, G., Botev, A., and Martens, J. (2022). Deep learning without shortcuts: Shaping the kernel with tailored rectifiers. In International Conference Learning Representations, volume 10, virtual.
12


Reviews:
[{"review_id": "3DBJOpJGtL", "reviewer": "Reviewer_kLTe", "summary": "NOTE: edited after author rebuttal and score has been updated.\n\nThis paper is related to input convex neural networks (ICNN). It analyzes the signal propagation through such a network and based on that proposes a new initialization scheme that allows the networks to be trained efficiently. It investigates the efficacy of the scheme in various ML tasks, including image benchmarks and drug discovery tasks.\n", "questions": "Minor spelling mistake 3.2 ”sufficiently good to the derive” - delete ”the”?\n", "limitations": "As mentioned in the paper, the signal propagation model assumes gaussian distributed preactivations, which is not always what happens in reality. This is well noted in the paper, and at least in the practical tasks considered in the empirical section, the derived initialization scheme still works well.\n\nAlthough the authors say that it would be possible to analyze skip connections, they do not perform the analysis. The reasoning that skip connections are not needed might not hold for deeper networks.\n", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The initialization scheme is well grounded in theory (although with some stated simplified gaussian assumption). \n\nThe method seems to work well in practice, at least in the tasks used in the paper.\n\nIt is experimentally shown that in contrary to previous work claims, ICNNs do not require skip connections in order to be trainable to good results.\n", "weaknesses": "The ICNN model family is intuitively less powerful than normal neural networks. The authors do not discuss this and the chosen experiments seem to be such that the performance degradation compared to normal neural networks does not happen. Could the authors discuss the limitations of the ICNN models, e.g., could imagenet level or GPT-style text understanding models be trained in the ICNN setting?\n\nThe abstract says that the new initialization allows for more efficient latent space exploration, but as far as I see, there is no comparison to other methods in the paper.\n\nFigure 1 shows preactivation distributions on the top, but the scale of the x axis is missing, please add that to the figure. It also seems that the activations start to concentrate around 0 as one moves deeper in the network. Is this the case and could the authors discuss this phenomenon and how this will affect very deep networks (e.g., 20 layers). \n\nIn the experimental sections, it seems that latent space exploration is the main use case enabled by the ICNN and the new init method? Could the authors discuss other use cases of ICNN and for example explain why they did not repeat the experiments from previous ICNN papers and show improvements stemming from their improved initialization? Also this use case seems very interesting, could the authors discuss pros and cons related to other methods that could be used to explore the latent space of this task? Also, are there some limitations of ICNN in this task or similar tasks, e.g., related to complexity or size of the network and the data? \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Limitations and Strengths of ICNNs\n\nYou are right that the ICNN family is strictly speaking less powerful than regular networks because they are constrained to have convex decision boundaries. \nNote that it is possible to construct a universal approximation theorem using theorem 1 from (Yuille and Rangarajan, 2001), which has been used in (Sankaranarayanan and Rengaswamy, 2022). \nOur manuscript might indeed give the (misleading) impression that ICNNs are direct competitors of regular networks. \nHowever, ICNNs are supposed to be used if the convexity provides additional benefits or is required for a method to work. \nWe provide the comparison with regular networks to emphasise that ICNNs are not necessarily harder to train than regular networks when using a principled (e.g. our proposed) initialisation.\nAs the reviewer points out, there is indeed no ICNN that reaches the quality of state-of-the-art regular networks at ImageNet or the text understanding quality of GPT. \nHowever, we do believe that the difficulty of training ICNNs is slowing research in this direction down and that our initialisation method might make ICNNs a more viable alternative.\n\nWe have added a few sentences to the experiment section to emphasise the goal of our comparisons with non-convex networks and to reduce this possible source of confusion.\n\n## Choice of Experiments\n\nIndeed, we show the latent space exploration for molecular generation as a particular example, in which ICNNs could be relevant. \nNotably, we do perform experiments from other ICNN papers, concretely the computer vision benchmarks were also included in (Sivaprasad et al., 2021). \nFor further examples where ICNNs are useful, we refer to our related work section.  \nIt might also be important to point out that our initialisation merely allows for more efficient training of ICNNs and is not the enabling component for these experiments.\n\nOur motivation behind these experiments  are some initial results by Nesterov et  al. (2022) who use ICNNs to “encourage” the latent space of an auto-encoder to build convex decision boundaries which allow efficient exploration of level sets. \nSeveral more methods for latent space exploration are compared in Du et al. (2023).\nThey introduce a new method, ChemSpaceE, for latent space exploration, which does not use ICNNs, and a way to evaluate these methods.\nHowever , we did not manage to reproduce their results and were therefore unable to include meaningful comparisons.\nFurthermore, these other methods for latent space exploration do not allow traversing level-sets, which is eventually the main feature provided by ICNNs and the property that we chose to highlight in our final experiment.\n\nWe will try to rephrase section 5.3 to resolve these sources of confusion:\n - the motivation for this particular experiment\n - the fact that ICNNs enable the level-set exploration, not our initialisation.\n\n### Additional References\n\nDu, Y., Liu, X., Shah, N. M., Liu, S., Zhang, J., & Zhou, B. (2022). ChemSpacE: Interpretable and Interactive Chemical Space Exploration. Transactions on Machine Learning Research.\n\n## Other Comments\n\nThe scale of the x-axis in figure 1 is implicit in the width of the bins, which is constant for all histograms in a sub-plot.\nHowever, we agree that this should have been explained and an explicit scale indication is the better way to communicate this information.\nThe activations tend to zero because any pre-activation that becomes negative is mapped to zero when using ReLU non-linearities.\nWe found that this can be alleviated by initialising the bias parameters with non-zero variance.\nDue to eq. (10), this also requires a different weight variance to keep the same propagation.\nOn the other hand, the instability of the fixed point (see appendix A.4) can lead to drift effects in the correlation, which effectively makes it hard to stabilise the propagation in very deep networks.\nFinding a setting where the correlation fixed point is actually stable would resolve this issue, but remains a problem for future work.\n\nWe have updated figure 1 to provide an indication of the scale of the x-axis for the final version (figure 1 in the rebuttal PDF).\nWe will also include a discussion about initialising the bias parameters with random samples drawn from a Gaussian distribution with variance $\\frac{1}{2}$.\nWe repeated the experiments from figure 5 in the appendix for this setting and find that the results are practically the same.\nFigure 2a in the rebuttal PDF shows the comparison between the initialisation with constant and random bias initialisation.\nThese results also suggest that the randomness in the biases might enable training even deeper networks.\n"}, {"review_id": "mjM9qDGxIH", "reviewer": "Reviewer_YxjJ", "summary": "This paper investigates the initialization for input-convex neural networks. They generalize the signal progagation theory by removing the assumption of centred weight distrubution. The experiments show that the proposed initiaization method is effective in a set of datasets.", "questions": "Please see weaknesses.", "limitations": "The authors havve adequately addressed the limitations, and there is no potential negative societal impact of their work.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1 This paper generalizes signal propagation theory by removing the assumption that weights are sampled from a centred distribution. This generalization is necessary for ICNNs.\n\n2 The experimetal results are solid and convincing. The experiment on the real-world drug discovery task is nice.", "weaknesses": "1 The theoretical contribution of this paper is limited. The theoretical results are derived  following the framework proposed in [1].\n \n2 It seems that the authors only consider the forward propagation of an initial ICNN. However, the backward propagation and the output diversity are also crucial for the initialization. \n\n\n\n\n\n\n[1] Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Theoretical contribution\n\nWe do not entirely understand why our generalisation of signal propagation theory and the derivation of the initialisation for ICNNs would be considered a _”limited”_ theoretical contribution.\nThe initialisation for ICNNs is indeed _limited_ to ICNNs.\nHowever, the generalised signal propagation theory could be useful outside of the scope of ICNNs as well.\nFor example, it could be used to study certain effects in regular networks due to the initial weights of a network not having exactly zero mean in practice.\nAfter all, our theory shows that feature correlation affects signal propagation as soon as the weights do not have zero mean.\n\nAs mentioned in the introduction and section 2 of our manuscript, we build on the framework presented in the thesis of Neal (1995), which was also the basis for the work from Poole et al. (2016).\nThe main contribution to signal propagation from Poole et al. (2016) was to include the correlation between samples in the analysis.\nThe work from Schoenholz et al. (2017) extends the propagation in (Poole et al., 2016) by including depth scales, dropout and the backward pass.\nFurther work extended the propagation for resnets (Yang et al., 2017) and convolutional layers (Xiao et al., 2018).\nThis information can also be found in the related work section of our manuscript.\nIn this context, we believe that our extension of signal propagation theory is not more _limited_ than any of these published works.\nFurthermore, we believe that our derivation of an initialisation for ICNNs is also not more _limited_ than e.g. the one that was recently derived for hypernetworks (Chang et al., 2020).\n\nFinally, our work does not include the correlation between samples, in contrast to (Poole et al., 2016; Schoenholz et al., 2017).\nWe show that the variance in a layer not only depends on the variance but also on the correlation between features in the previous layer when weights have non-zero mean.\nAs a result, the propagation of correlation between features interferes with the propagation of variance, which is arguably more complex than studying independent signals as in (Poole et al., 2016; Schoenholz et al., 2017).\nNote that although the expressions are very similar, the correlation between samples and the correlation between features are two different things.\nIn this sense, we would argue that our work is not _limited_ to a simple derivation of (Poole et al., 2016) or (Schoenholz et al., 2017), but rather provides a different perspective into signal propagation theory.\n\n## Backward Analysis\n\nWe agree that incorporating the analysis of the backward pass to derive an initialisation is generally desirable.\nTherefore, we decided to include the generalised signal propagation of the backward pass in the appendix.\nThe mean and variance propagation of the deltas in backpropagation (using similar assumptions as Schoenholz et al. (2017) is given by\n$$\\begin{align*}\n\t\\mathbb{E}[\\delta_j^{-}] &= M \\mu_w \\mathbb{E}[\\phi'(s_1^{-})] \\mathbb{E}[\\delta_1] \\\\\\\\\n\t\\mathbb{E}[\\delta_i^{-} \\delta_j^{-}] &= \\updelta_{ij} M \\sigma_w^2 \\mathbb{E}\\bigl[\\phi'(s_1^{-})^2\\bigr] \\mathbb{E}\\bigl[\\delta_1^2\\bigr] + \\mu_w^2 \\mathbb{E}[\\phi'(s_i^{-}) \\phi'(s_j^{-})] \\sum_{k,k'} \\mathbb{E}[\\delta_k \\delta_{k'}],\n\\end{align*}$$\nWhere $\\updelta_{ij}$ is the Kronecker delta and $\\delta_i = \\frac{\\partial L}{\\partial s_i}$.\nNote that, similar to the forward pass, the covariance between entries in the delta vectors interferes with the variance propagation.\nThe kernel for the derivative of $\\operatorname{LReLU}$ with parameter $\\alpha$ is given by\n$$\\mathbb{E}[\\operatorname{LReLU}'(s_1 \\mathbin{;} \\alpha) \\operatorname{LReLU}'(s_2 \\mathbin{;} \\alpha)] = (1 - \\alpha)^2 \\frac{1}{2 \\pi} \\arccos(-\\rho) + \\alpha.$$\nThis should make it possible to derive initialisations that also incorporate the backward pass, e.g. using appraoches from (Glorot et al., 2010) or (Defazio et al., 2021).\n\n## Experiments\n\nWe thank you for the positive assessment of our experimental section.\nWe also think that the latent-space exploration of molecules is a particularly illustrative example of the usability of ICNNs."}, {"review_id": "JWAsPOVJ0M", "reviewer": "Reviewer_f5Cq", "summary": "The proposed method aims to solve the initialisation problem in an Input convex Neural Network (ICNN), where the weights are required to the non-negative. The commonly applied approach, setting the negative entries stamped from zero-mean Gaussian as zero, leads to varies the desired mean. By analysing the signal propagation in the Neural Networks, the authors are able to sample the weights with expected mean and variance. The method is evaluated on several tasks and compared with the ICNN without the proposed initialisation method with some promising results. ", "questions": "1. In Figure 2, Non-convex and ICNN have lower training loss than other baselines at iteration 0 which, I believe, indicates the initialisation can the author explain why it happens?  \n2. Can we treat $\\rho_*$ as a hyperparameter? If so how it affects training and generalisation of the ICNN models? these two questions are not discussed.  ", "limitations": "See the sections above. ", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The problem solved in the submission is important in the ICNN setting.  \n2. The derivation of the method is concrete. ", "weaknesses": "Some details are not well explained and are sometimes confusing. \n1. $\\rho_*$ is defined $\\frac{1}{\\sigma^2_*}Cov[s_1^-, s_2^-]$ in line 178 and the author claim $\\rho_*$ is independed on $\\frac{1}{\\sigma^2}$ in line 184. And then  \\rho_* is set as $\\frac{1}{2}$ heuristically. \n\nThe experiments are not carefully designed to support the submission and I am not sure some experiments are sufficiently conducted. \n1. In Figure 2, the loss of ICNN does not change during the training process in the MNIST setting but in CIFAR10 and CIFAR100 there is no such phenomenon. \n2. In Figure 3, drawing the conclusion that with the proposed initialisation the training dynamic is more stable is not concrete via eyeball comparison. Some quantities are needed. \n3. And according to my experience of training non-convex NN on cifar10, the average testing loss does not really show a dramatic increase after some training iterations and non-convex NN has the highest test loss and accuracy which are unusual. And the rest of the learning curve is more reasonable with low test loss and corresponding high accuracy. \n4. It can be noticed that with skip on ICNN the performance of ICNN improves, I think applying the proposed initialisation to ICNN with skip will give good support to the submission. \n\nAdditional question, which does not affect my score. Since in all the settings in the submission, the non-convex NN has the best performance compared with other ICNN-based methods, so is there any scenario for ICNN or why ICNN is essential in the first place? \n\nNotations: \nIn some equations the $s_i$, $s_j$ are mixed used with $s_1$, $s_2$, for example, the one below line 172 and the one below 175. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank you for your feedback. We hope to address all of your points in detail soon!\n\nFor now, we would like to ask for some further clarification on what quantities the reviewer would like to see in the comparison.\nThis would help us to better address your concerns in our detailed response.\n\nThanks in advance!\n\n---\n\nWe thank the reviewer for their critical feedback and apologise for details that were confusing or unclear. \nWe try to clarify these details better in the answers below and the final version.\nWe hope that you will take the time to consider our rebuttal and are willing to update your score after possible further discussions.\n\n## Weaknesses\n\nConcerning the unclarity:\n\n1. We agree that the statement might be confusing, but it is inherent to the definition of correlation.\nBecause the correlation is defined in terms of the variance they can not be independent.\nOur statement about eq. (11) and (12) being independent aims to make clear that once we have computed the correlation, we do not explicitly need the variance to derive our initialisation.\nWe reformulated line 184 to avoid confusion.\n\n   The choice for setting $\\rho_* = \\frac{1}{2}$ can indeed be considered as a heuristic.\n   Our main motivation was to get a closed-form result for $\\arccos(\\rho_*)$, but other values are possible as well (more details below).\n   We will add a section in the appendix with these details in the final revision.\n\nConcerning the experiments:\n\n1. Note that the loss for ICNN does actually go down during the first hundred update steps, but we agree that this is barely visible in the plot.\nThe CIFAR-10 and CIFAR-100 plots for ICNNs without skip-connections or initialisation show the same behaviour: the training loss goes down initially, but stops improving at some point.\nWe suspect that this happens because the network attempts to reduce the activation strength of the network by pushing weights and biases down.\nIn the example of the MNIST models, the biases in the first unconstrained layer soon become all negative.\nAs a result, the ReLU activations, which are the inputs to the first constrained layer, are all zero and the network can not learn a function from the inputs.\nThe skip-connections alleviate these issues by bypassing these dead layers.\nOur initialisation simply provides a better starting point where activations do not need to be reduced to improve the error early in training.\nWe will try to include a more elaborate discussion in the final version of the paper.\n2. We could not find the point where we conclude that training dynamics are more stable in the context of Figure 3.\nFigure 3 aims to show that our initialisation does not only affect the empirical error but also translates to generalisation performance.\nWe agree that there is little to no difference for the MNIST experiments and we also acknowledge that in our manuscript (line 283).\nHowever, we would argue that for the CIFAR10 experiments, our method (ICNN + init) clearly leads to faster learning compared to other ICNNs.\nThe quantitative accuracies for the validation runs can be found in Table 3 in the appendix.\n3. The increase in the test loss in Figure 3 can be explained by overfitting on the training and/or validation data.\nNote that the optimal models were chosen based on validation accuracy, as indicated by Table 3 in the appendix.\nBecause the loss is only a proxy for accuracy, the optimal model might indeed be in an overfitting regime in terms of test loss.\n\n   By revisiting Table 3, we realised that Figure 3 does not depict the early stopping.\n   We updated the figure for the final version (Figure 3 in the rebuttal PDF).\n4. One of our contributions is to show that ICNNs do not need skip-connections if they are initialised in a principled manner.\nThis indicates that skip-connections mainly help in making ICNNs trainable.\nWe did (accidentally) run the Tox21 experiments with both skip-connections and our initialisation and did not observe any substantial improvements over using only our initialisation.\nHowever, the skip-connections do effectively modify the signal propagation in a non-trivial way.\nIt would require additional analysis to obtain a principled approach for initialising networks with skip-connections.\nWe also refer to a similar reply we gave to reviewer Ryn9 with more references in this context.\n\nConcerning your additional questions:\n\n - The performance of ICNNs should not be compared directly with regular networks.\nAs indicated by reviewer kLTe, ICNNs have theoretically less capacity.\nICNNs have the unique property that they are convex.\nThis can be useful or is even necessary in various settings (e.g. energy-based models, optimal transport, level-set exploration, …) that we describe in our related work section.\n - The notation on lines 172 and 175 is on purpose.\nWe aim to compute the (co)variance for arbitrary $s_i$ and/or $s_j$ on the left-hand side.\nUnder our assumptions, the (co)variance turns out to be independent of the index $i$ or $j$.\nWe explicitly use $s_1$ and $s_2$ on the right-hand side to emphasise this independence.\n\n## Questions\n\n1. The ICNNs have only non-negative weights.\nDue to the ReLU activation function, also activations will be positive.\nComputing these dot products typically leads to numerically large values.\nThis is also why the loss will typically be large.\nOur initialisation counters these effects mainly by initialising the bias parameters with negative values.\n2. $\\rho_*$ can indeed be treated as a hyper-parameter.\nBecause of the tendency of the correlation to grow as the network grows deeper, lower $\\rho_*$ values will typically make it possible to train deeper networks.\nWe ran ablation experiments on the choice for $\\rho_*$ to obtain figure 2b in the rebuttal PDF. \nAs expected, lower values for $\\rho_*$ can enable deeper networks, but overall, performance is very similar.\nWe will include these results in the appendix of the final version."}, {"review_id": "NXtlLCtzlX", "reviewer": "Reviewer_Ryn9", "summary": "The paper discusses a principled weight initialization strategy for Input-Convex Neural Networks (ICNNs). The authors propose a new theory that generalizes signal propagation theory to include weights without zero mean, and derive a principled initialization strategy for ICNNs from this theory. They demonstrate the effectiveness of their initialization strategy through empirical experiments and apply ICNNs in a real-world drug-discovery setting.\n", "questions": "N/A", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper proposes a new theory that generalizes signal propagation theory to include weights without zero mean, which is a significant contribution to the field.\n2. The authors derive a principled initialization strategy for ICNNs from their new theory, which improves learning and generalization in ICNNs.", "weaknesses": "1. The proposed approach cannot be applied to networks with skip connections, limiting its usage in real-world models.\n2. The experimental section is not comprehensive and there is a lack of ablation study", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank you for your feedback. We hope to address all of your points in detail soon!\n\nFor now, we would like to ask for some further clarification on what ablation studies are missing.\nThis would help us to better address your concerns in our detailed response.\n\nThanks in advance!\n\n---\n\nWe thank the reviewer for their positive and encouraging feedback.  \nWe hope that we can shed light on the mentioned weaknesses. \n\n##  Skip connections\n\nOur proposed initialisation scheme was developed to enable faster training in ICNNs.\nBecause skip-connections have a similar purpose, we would argue that our initialisation can be a replacement for skip-connections. \nPrior to our work, CNNs could not be trained without skip-connections. \nWith our principled initialisation, we make it possible to train ICNNs without skip-connections.\nFurthermore, our initialisation appears to enable better results than ICNNs with skip-connections, which have more parameters.\n\nThere is also a strong trend to replace skip-connections in regular networks.\nFor example in (Zhang et al., 2022) it is shown that regular deep networks can be trained to the same performance as ResNets by carefully controlling signal propagation. \nAlso, note that skip-connections in ICNNs usually connect layers with the input to the network and not block-wise, as in residual networks. \nTherefore, ICNNs used in prior work have a strong tendency towards dynamics that are similar to single-layer networks and suffer from the feature reuse problem (Zagoruyko et al., 2016). \nThe ICNNs that we trained in this work have to develop feature hierarchies and are therefore a step forward for researchers working with ICNNs.\n\nFurthermore, we emphasise that our method can be applied to networks with skip-connections in practice. \nIt is just the theory that is difficult to derive because of the dependency between the inputs and outputs of the residual branch.\n\n### Additional References\n - Zagoruyko, S., & Komodakis, N. (2016). Wide Residual Networks. Proceedings of the British Machine Vision Conference 2016, 87.1-87.12. https://doi.org/10.5244/C.30.87\n - Zhang, G., Botev, A., & Martens, J. (2022). Deep Learning without Shortcuts: Shaping the Kernel with Tailored Rectifiers. International Conference on Learning Representations 10. https://openreview.net/forum?id=U0k7XNTiFEq\n\n## Ablation study\n\nFirst of all, we would like to point out that our main experiments are already an ablation study.\nWe start from ICNNs with skip-connections (the most common form of ICNN at the time of writing) and establish the baseline performance. \nThe first ablation is obtained by removing the skip-connections which results in networks that become practically impossible to train.\nFinally, we add our principled initialisation to the networks without skip-connections and find that the initialisation improves not only the networks without skip-connections, but also the networks with skip-connections.\nUnfortunately, we were unable to deduct which ablations you would have liked to see in addition.\n\nWe included learning curves for different choices of $\\rho_*$ as suggested by reviewer f5Cq (figure 2c in the rebuttal PDF).\nWe also investigated what happens if we choose to initialise the biases with $\\sigma_b > 0$ (figure 2a in the rebuttal PDF).\nFinally, we added an experiment to see how important the effect of the bias shift is by only using eq. (8) for initialising the bias parameters.\nThe weight parameters were initialised using an initialisation scheme for regular networks.\nThe results of this experiment can be found in figure 2b of the rebuttal PDF.\nAs suspected, only initialising the bias parameter already leads to good performance, but results tend to be better/more consistent when including the weights for the initialisation."}]

Author rebuttal:
We thank the reviewers for their feedback and suggestions.
The comments and concerns of each reviewer will be addressed individually and the paper will be updated correspondingly.

We did observe a common (potential) misunderstanding about our contribution in the reviews.
Some reviewers claim in their summary that we _“analyse”_ the signal propagation to derive an initialisation.
Although _analyse_ captures the detailed study of signal propagation, we believe it does not quite capture the value that we believe to have added to the traditional theory.
Therefore, we would like to emphasise that we do not just use signal propagation theory as a tool to derive an initialisation.
Instead, we generalise the existing signal propagation theory to allow initial weights with non-zero means.
Only then do we have the signal propagation tools to derive an initialisation for ICNNs.
Possibly this was already clear for the reviewers, but we want to avoid any misunderstandings.
A detailed discussion can be found in the reply to reviewer YxjJ.

We look forward to further feedback during the discussion phase.

Meta-review: While its rating puts it squarely in the borderline zone, this paper makes a sound and sufficiently novel (if somewhat narrow) contribution. It appears that some of the reviewers with lower ratings might have misunderstood core aspects of the paper, which the authors addressed in a strong and thoughtful rebuttal that unfortunately went unanswered. Although additional experiments would have been useful, I side with reviewer kLTe in accepting the existing results as sufficient to validate the main claim of the paper: that the proposed framework provides a principled, efficient, and useful approach to initialize ICNNs.
Decision: Accept (poster)