Input convex neural network (ICNN)

Overview

An input convex neural network (ICNN) is a feed-forward network that is convex in its input. Each layer adds an unconstrained affine map of the input to the previous hidden state. The network stays convex if the hidden-to-hidden weights are non-negative, the first activation is convex, and every later activation is convex and non-decreasing 1908.10962 §3.1. The input enters at every layer. The paper cites Chen et al. for the claim that ICNNs can approximate any convex function on a compact set 1908.10962 §3.1.

Only one card lists this slug. The page therefore describes a single use, the pair of potentials in a semi-dual optimal transport problem, and the choices the paper makes for it. Other uses of ICNNs are not covered here.

Variants

Potentials for quadratic-cost optimal transport (minimax)

Two ICNNs, f and g, are trained against each other in a minimax game. The transport map is the gradient of g 1908.10962 §3. The source distribution must have a density for the optimal map to be unique and given by a gradient 1908.10962 §3, thm. 3.1. The paper reformulates the squared 2-Wasserstein distance as a supremum over convex f of an infimum over convex g, plus a constant 1908.10962 §3, thm. 3.2, eq. 5. The g-objective adds a penalty on negative hidden-to-hidden weights, which relaxes the convexity constraint on g 1908.10962 §4, eq. 8. The paper contrasts this with a conjugate-based method that needs a convex solve for each sample 1908.10962 §3.1. Card: Optimal transport mapping via input convex neural networks.

On a 2-D checkerboard and 8-Gaussian setting, the paper reports that its map does not leave points in the gaps between modes, unlike the compared methods 1908.10962 §4.1, fig. 3. This is a qualitative result. On Gaussian pairs in 784 dimensions, the author-reported squared error in the transported mean is 0.19, 13.95 and 29.05 for mean shifts of 1, 5 and 10, so the absolute error grows with the shift. The paper reports a ± value next to each number but does not define it 1908.10962 tab. 1.

When to use which

The only variant is the minimax setup above, so the table compares the choice of whether to use it at all.

Implementation sketch

The core is one ICNN forward pass and an alternating minimax loop. The pseudocode below is written in our own words. Batch sizes, step counts and the exact update order are in the cards' tables and in the pitfalls.

# Forward pass of an ICNN  [arxiv:1908.10962 §3.1]
# Input: x (M, d). A_l, b_l unconstrained. W_l >= 0 for l >= 1.
# Output: scalar value per sample (M, 1)

def icnn(x):
    # Layer 0 has no W term; the activation must be convex  [§3.1]
    hidden = act0(x @ A_0 + b_0)                       # (M, m)
    for l in range(1, L):
        # The hidden-to-hidden path uses non-negative weights  [§3.1]
        hidden = act(x @ A_l + b_l + hidden @ W_l)     # (M, m)
    # The output also passes through its activation  [§3.1]
    return act(x @ A_L + b_L + hidden @ W_L)           # (M, 1)

# Minimax loop with the transport map as grad g  [alg. 1] [§3]
def train(Q, P, M, K, T):
    for t in range(T):
        # K updates of g, each with a fresh source batch  [alg. 1]
        for k in range(K):
            Y = sample(Q, M)                           # (M, d)
            maps = grad_wrt_input(g, Y)                # (M, d)
            # Penalty on negative weights of g  [§4, eq. 8]
            loss_g = mean(f(maps)) - mean(sum(Y * maps, 1))
            g = adam_step(g, loss_g + penalty(g))
        # One update of f, then clip its hidden weights  [alg. 1]
        X = sample(P, M)                               # (M, d)
        loss_f = mean(f(X)) - mean(f(maps))
        f = adam_step(f, loss_f)
        f = clip_hidden_weights(f, minimum=0)

Contradictions

None known. Only one card exists, so no two papers disagree. The card's Paper vs code section lists places where the authors' code differs from the paper, for example the penalty scale and whether the source batch is redrawn in each inner step 1908.10962 §4, alg. 1. See the card for the full list.

Pitfalls

• The g-step needs second-order autodiff. The g-objective depends on grad g, so training differentiates through a gradient [arxiv:1908.10962 §4, alg. 1].
• Positivity is enforced in two ways. f's hidden-to-hidden weights are clipped at zero after each update, while g receives a penalty instead 1908.10962 alg. 1, §4, eq. 8.
• Keep the output activation. The paper's layer equation applies the activation to the last layer too 1908.10962 §3.1. Leaving the output unactivated changes the function class.
• The loss is W_2 squared, with the factor one half. The paper's quantity is the squared distance with a factor one half 1908.10962 §2, eq. 2.
• Which batch the f-step uses is not stated. Algorithm 1 does not name the batches for the f-update 1908.10962 alg. 1.
• Initialisation is not stated. The paper gives no weight initialisation for its experiments 1908.10962 app. C.1, app. C.2.
• The leaky-ReLU slope for high dimensions is not stated. The paper gives beta = 0.2 for the 2-D case only 1908.10962 app. C.1, app. C.2.
• Check with an exact solver on small batches. The authors' repository uses an exact discrete solver on uniform samples to check the estimate. The card records the file reference W2-minimax-tf.py.

Papers

• Optimal transport mapping via input convex neural networks: the minimax, semi-dual ICNN method for quadratic-cost optimal transport from samples, with 2-D and 784-D experiments.