Title: Principled Weight Initialisation for Input-Convex Neural Networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: initialization

signal propagation

input-convex networks
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

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.

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 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 inputconvex 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.,

pre-activation

distributions

weight signs

<

0

or

0

feature

correlations

0

5

0

5

0

5

(a) traditional

0

5

5

8

51 80

(b) ICNN

0

5

0

20

0

40

(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 skipconnections 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 skipconnections. This confirms the hypothesis that good signal propagation can replace skip-connections (Martens et al., 2021; Zhang et al., 2022).

Contributions. Our contributions 1 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 fullyconnected 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).

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 : R N → R M , with some activation function ϕ : R → R that is applied element-wise to vectors. We study the propagation of pre-activations 2 (cf. Saxe et al., 2014; He et al., 2015; Poole et al., 2016) throughout the network:

Here, W ∈ R M × N and b ∈ 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 -∈ 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 ∼ 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 ∀ 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 E [ w ij ] = 0 , and Var[ w ij ] = σ 2 w , the first two moments of the pre-activations are given by

̸

Note that the variance is independent of the index, i , in the pre-activation vector and thus ∀ k : E [ ϕ ( s -k ) 2 ] = E [ ϕ ( s -1 ) 2 ] . Moreover, it can be shown (see Appendix A) that ∀ i = j : E [ s i s 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. ∀ i, j : E [ s 2 i ] = E [ s -j 2 ] = σ 2 ∗ . Plugging this into Eq. (3), we obtain the fixed-point equation

Here, varprop ϕ is a function that models the propagation of variance through the activation function, ϕ , assuming zero mean inputs. E.g. varprop ReLU ( σ 2 ) = 1 2 σ 2 (He et al., 2015) or varprop tanh ( σ 2 ) ≈ σ 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 nonnegative weights in ICNNs. However, Eq. (4) heavily relies on the assumption that the weights are centred, i.e. µ w = 0 . This assumption is impossible to satisfy for the non-negative weights in ICNNs, unless all weights are set to zero. If µ w = 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.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 E [ w ij ] = µ w = 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 E [ b i ] = µ b and variance Var[ b i ] = σ 2 b . By re-evaluating the expectations in Eq. (2) and Eq.(3), we arrive at the following results:

where δ 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 ∀ k : E [ ϕ ( s -k ) n ] = E [ ϕ ( s -1 ) n ] .

̸

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, E [ s 2 i ] , 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, Cov[ s i , s j ] = E i = j [ s i s j ] -E [ s i ] 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:

with ( s 1 , s 2 ) ∼ N ( (0 , 0) , ( 1 ρ ρ 1 ) σ 2 ) and correlation ρ . Note that this expression implicitly de- scribes the propagation of the squared mean ( ρ = 0) and the second raw moment ( ρ = 1 ). We will assume ϕ = 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 µ w = 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 µ w , σ 2 w , µ b and σ 2 b 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 ◦ 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. softplus (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 skipconnections 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 preactivations have zero mean. Therefore, we look for a configuration that makes the mean (Eq. 5) of the pre-activations zero. Because we cannot enforce µ w = 0 , we make use of the bias parameters to obtain centred pre-activations:

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

̸

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

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:

̸

with ρ ∗ = 1 σ 2 ∗ Cov[ s -1 , s -2 ] = 1 σ 2 ∗ Cov i = j [ s i , s j ] and σ 2 ∗ = Var[ s -1 ] = 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 σ 2 b , σ 2 w and µ 2 w . Because this system is over-parameterised, we choose to set σ 2 b = 0 . An analysis for σ 2 b > 0 can be found in Appendix C.3. This leads to the following initialisation parameters for ICNNs:

Note that these solutions only depend on the correlation, ρ ∗ , and not on the variance, σ 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, ρ = Cov[ s 1 ,s 2 ] Var[ s 1 ] , not all solutions are admissible. For example, uncorrelated features ( ρ = 0 ) are only possible if µ w = 0 . This means that features in ICNNs must have non-zero correlation. On the other hand, perfect correlation ( ρ = 1 ) would mean that all features point in the same direction, which is not desirable. Therefore, we choose ρ ∗ = 1 2 as a compromise to obtain the following initialisation parameters for our experiments in Section 5:

We refer to appendix C.4 for a discussion on different choices for ρ ∗ .

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

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 ˜ w ij ∼ N ( ˜ µ w , ˜ σ w 2 ) to compute w ij = exp ( ˜ 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.

4 Related Work

Input-convex neural networks. Input-Convex Neural Network (ICNN) were originally designed in the context of 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.e. 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 to use a 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 an inherent part of the forward pass. Another alternative is to use regularisation to impose a soft non-negativity constraint on the weights (Makkuva et al., 2020). Similar to (Sivaprasad et al., 2021), we directly train an ICNN, but we do not observe the same generalisation benefits. On the other hand, Sankaranarayanan and Rengaswamy (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 a connection 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 modern 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 adaptations to account for weights with non-zero mean. There is also a 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, we 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 generalisation in ICNNs on multiple permuted image datasets. We also include non-convex networks in these experiments to illustrate that ICNNs with our principled initialisation can be trained as well as regular networks. However, we would like to stress that non-convex networks are expected to outperform 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 a tool for more controlled latent space exploration of a molecule auto-encoder system. This illustrates 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., 1994), CIFAR10 and CIFAR100 (Krizhevsky, 2009), to which we refer as permuted image benchmarks (cf. Goodfellow et al., 2014).

In our comparison, we consider four different types of methods: (1) a classical non-convex network without skip-connections and default initialisation (cf. He et al., 2015). (2) an ICNN, 'ICNN', without 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 is implemented by setting negative weights to zero after every update (cf. Amos et al., 2017). Note that skip-connections in ICNNs introduce additional parameters to allow for additional representational power (cf. Amos et al., 2017).

Training dynamics. Settings. 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 was obtained by manually tuning on the non-convex baseline. Results. We found that ICNNs with our principled initialisation exhibit similar learning behaviour as the non-convex baseline, while ICNNs without our initialisation strategy can not decrease the loss to the level of the baseline on CIFAR10 and CIFAR100. Figure 2 shows the learning curves for our different methods on three permuted image benchmarks.

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 performance this time. To this end we perform a grid-search to find the hyper-parameters that attain the

MNIST

CIFAR10

CIFAR100

101

102

ICNN

ICNN + skip

non-convex

ICNN + init

101

103

3 ×100

2 ×100

103

4 ×100

5 × 100

102

WWAME

training loss

100

10-3

100

102

3 ×100

10-1

101

101

10-2

100

10-3

0

2000

4000

6000

8000

0

1000

2000

3000

4000

5000

6000

7000

0

1000

2000

3000

4000

5000

6000

7000

100

1.00

6 ×100

0.60

non-convex

non-convex

ICNN

ICNN

ICNN + skip

ICNN + init

0.98

4 ×100

ICNN + skip

ICNN + init

0.55

0.50

average loss

0.96

accuracy

3×100

accuracy

0.45

0.94

2 × 100

0.40

101

0.92

0.35

0

5

10

15

20

25

0

5

10

15

20

25

0.90

0

5

10

15

20

25

0

5

10

15

20

25

0.30

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, L 2 -regularisation and whitening pre-processing transforms (details in Appendix C.2). Results. On the MNIST dataset, all variants reach similar accuracy values, which we attribute to the simplicity of the prediction problem. However, on CIFAR10 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, we 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 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 a 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 1 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.

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 in drug discovery to keep one molecular property, such as 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 in the direction of another randomly chosen target molecules. Since the input space of the ICNN that predicts toxicity (see above) possesses a decoder model from (Winter et al., 2019), the numeric representations of the molecules on the level set can be decoded into molecular structures. The obtained trajectories in Figure 4 demonstrate that the predicted toxicity of the molecules remain constant, while different QEDscores, 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 a stable propagation of signals. We empirically verified that this steady propagation leads to improved learning. Finally, we showed that ICNNs can be trained to a 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 of 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 proficient 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 included a signal propagation analysis for back-propagation 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 nonnegative weights.

Acknowledgments and Disclosure of Funding

We thank Sepp Hochreiter for pointing us in the direction of input-convex networks and their interesting 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-6YOU-212), DeepFlood (LIT-2019-8-YOU-213), Medical Cognitive Computing Center (MC3), INCONTROL-RL (FFG-881064), PRIMAL (FFG-873979), S3AI (FFG-872172), DL for GranularFlow (FFG-871302), EPILEPSIA (FFG-892171), AIRI FG 9-N (FWF-36284, FWF-36235), AI4GreenHeatingGrids(FFG- 899943), INTEGRATE (FFG-892418), ELISE (H2020-ICT-20193 ID: 951847), Stars4Waters (HORIZON-CL6-2021-CLIMATE-01-01). We thank Audi.JKU 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 Competence 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 in 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 in 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 performance gap in unnormalized ResNets. In International Conference on Learning Representations , volume 9, virtual. arXiv: 2101.08692.

Chang, O., Flokas, L., and Lipson, H. (2020). Principled weight initialization for hypernetworks. In International Conference on 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). In 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 a Dual View on Expressivity. In Advances in Neural Information Processing Systems , volume 29, pages 2253-2261, Barcelona, Spain. Curran Associates, Inc.

Defazio, A. and Bottou, L. (2021). Beyond Folklore: A Scaling Calculus for the Design and Initialization 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., Belgrave, D., and Cho, K., editors, Advances in Neural Information Processing Systems , volume 35.

Goodfellow, I. J., Mirza, M., Da, X., Courville, A. C., and Bengio, Y. (2014). An Empirical Investigation of Catastrophic Forgeting in Gradient-Based Neural Networks. In International Conference on Learning Representations , volume 2.

He, K., Zhang, X., Ren, S., and Sun, J. (2015). Delving deep into rectifiers: Surpassing human-level performance on 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, Technical University Munich, Munich, Germany.

Hoedt, P.-J., Hochreiter, S., and Klambauer, G. (2018). Characterising activation functions by their backward dynamics around forward fixed points. NeurIPS workshop on Critiquing and Correcting Trends in Machine Learning.

Hoedt, P.-J., Hochreiter, S., and Klambauer, G. (2022). Normalization is dead, long live normalization! In ICLR Blog Track . https://iclr-blog-track.github.io/2022/03/25/unnormalized-resnets/.

Huang, R., Xia, M., Nguyen, D.-T., Zhao, T., Sakamuru, S., Zhao, J., Shahane, S. A., Rossoshek, A., and Simeonov, A. (2016). Tox21Challenge to Build Predictive Models of Nuclear Receptor and Stress Response Pathways as Mediated by Exposure to Environmental Chemicals and Drugs. Frontiers in Environmental Science , 3.

Klambauer, G., Unterthiner, T., Mayr, A., and Hochreiter, S. (2017). Self-normalizing neural networks. 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, A. 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 on Machine Learning , volume 119, pages 6672-6681. PMLR.

Martens, J., Ballard, A., Desjardins, G., Swirszcz, G., Dalibard, V., Sohl-Dickstein, J., and Schoenholz, S. S. (2021). Rapid training of deep neural networks without skip connections or normalization layers using Deep Kernel Shaping. arXiv: 2110.01765.

Mayr, A., Klambauer, G., Unterthiner, T., and Hochreiter, S. (2016). Deeptox: toxicity prediction using deep learning. Frontiers in Environmental Science , 3:80.

Mishkin, D. and Matas, J. (2016). All you need is a good init. In International Conference on Learning Representations , volume 4.

Nair, V. and Hinton, G. E. (2010). Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th International Conference on Machine Learning , pages 807-814, Madison, 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. Neurocomputing , 495:153-168.

Saxe, A., McClelland, J., and Ganguli, S. (2014). Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In International Conference on Learning Representations , volume 2.

Schoenholz, S. S., Gilmer, J., Ganguli, S., and Sohl-Dickstein, J. (2017). Deep information propagation. 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 lstm: A convex approach for fast Lyapunov-based model predictive control. arXiv preprint arXiv:2311.07202 .

Winter, R., Montanari, F., Noé, F., and Clevert, D.-A. (2019). Learning continuous and datadriven molecular descriptors by translating equivalent chemical representations. Chemical Science , 10(6):1692-1701. Publisher: The Royal Society of Chemistry.

Xiao, L., Bahri, Y., Sohl-Dickstein, J., Schoenholz, S., and Pennington, J. (2018). Dynamical isometry and a mean field theory of cnns: How to 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 on Learning Representations , volume 10, virtual.

Peer reviews:
comments: 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.

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.
2. 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.
2. The experimental section is not comprehensive and there is a lack of ablation study

QUESTIONS:
N/A

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-28

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. The problem solved in the submission is important in the ICNN setting.  \n2. The derivation of the method is concrete.\n\nWEAKNESSES:\nSome 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nSee the sections above.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-02"}

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

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.

2 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].
 
2 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. 






[1] Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017

QUESTIONS:
Please see weaknesses.

LIMITATIONS:
The authors havve adequately addressed the limitations, and there is no potential negative societal impact of their work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nNOTE: 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\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nThe 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\nQUESTIONS:\nMinor spelling mistake 3.2 ”sufficiently good to the derive” - delete ”the”?\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-24"}