Input Convex Neural Networks with Flux.jl

This tutorial shows how to embed an input convex neural network (ICNN) model from Flux.jl into JuMP.

Required packages

This tutorial requires the following packages:

using JuMP
import Flux
import HiGHS
import Ipopt
import MathOptAI
import Plots
import Random
import SCS

The following custom layer can be used to build ICNNs. This layer has two forward methods. One that takes a single input and the other takes a Tuple. They both return the result of the forward pass as well as the original input.

struct InputConvex{T,F}
    weight_x::Matrix{T}
    weight_z::Matrix{T}
    bias::Vector{T}
    σ::F
end

Flux.@layer(InputConvex, trainable = (weight_x, weight_z, bias))

function InputConvex(
    ((in_z, in_x), out)::Pair{Tuple{Int,Int},Int},
    σ = identity;
    init = Flux.glorot_uniform,
)
    return InputConvex(init(out, in_x), init(out, in_z), init(out), σ)
end

function (c::InputConvex)(x::AbstractVector)
    return c.σ.(c.weight_x * x .+ c.bias), x
end

function (c::InputConvex)((z, x)::Tuple)
    return c.σ.(Flux.softplus.(c.weight_z) * z .+ c.weight_x * x .+ c.bias), x
end

function Base.show(io::IO, l::InputConvex)
    m, n = size(l.weight_x)
    print(io, "InputConvex((", size(l.weight_z, 2), ", $m) => $n")
    if l.σ != identity
        print(io, ", ", l.σ)
    end
    if l.bias == false
        print(io, "; bias=false")
    end
    print(io, ")")
    return
end

Here's an example:

layer = InputConvex((8, 8) => 2, Flux.relu)

InputConvex((8, 2) => 8, relu)  # 34 parameters

layer(rand(8))

([1.21138549550919, 2.7904586216549747], [0.05210711744747376, 0.5840484444114491, 0.5708444178206272, 0.8992055628745603, 0.4989080944156268, 0.5094963690976293, 0.9152891073151073, 0.684746000696546])

Next, we define a custom Chain to build the ICNN.

struct InputConvexChain{T<:Flux.Chain}
    chain::T
end

InputConvexChain(layers...) = InputConvexChain(Flux.Chain(layers))

(model::InputConvexChain)(x) = first(model.chain(x))

function Base.show(io::IO, l::InputConvexChain)
    println(io, "InputConvexChain(")
    println.(io, "\t", l.chain)
    println(io, ")")
    return
end

Here's an example:

chain = InputConvexChain(
    InputConvex((8, 8) => 2, Flux.relu),
    InputConvex((2, 8) => 1, Flux.relu),
)

InputConvexChain(
	InputConvex((8, 2) => 8, relu)
	InputConvex((2, 1) => 8, relu)
)

chain(rand(8))

1-element Vector{Float64}:
 0.0

Building the Predictor

We need to implement build_predictor and add_predictor for InputConvexChain in order to be able to embed this network into JuMP.

struct InputConvexChainPredictor <: MathOptAI.AbstractPredictor
    p::MathOptAI.Pipeline
end

function MathOptAI.build_predictor(
    predictor::InputConvexChain;
    config::Dict = Dict{Any,Any}(),
    kwargs...,
)
    (layer1, layers) = Iterators.peel(predictor.chain)
    p = MathOptAI.Pipeline(
        MathOptAI.Affine(layer1.weight_x, layer1.bias),
        MathOptAI.build_predictor(layer1.σ; config),
    )
    for layer in layers
        weights = hcat(Flux.softplus(layer.weight_z), layer.weight_x)
        push!(p.layers, MathOptAI.Affine(weights, layer.bias))
        push!(p.layers, MathOptAI.build_predictor(layer.σ; config))
    end
    return InputConvexChainPredictor(p)
end

function MathOptAI.add_predictor(
    model::JuMP.AbstractModel,
    predictor::InputConvexChainPredictor,
    x::Vector;
    kwargs...,
)
    layers = predictor.p.layers
    z, inner = MathOptAI.add_predictor(model, first(layers), x)
    formulation = MathOptAI.PipelineFormulation(predictor, Any[inner])
    for layer in layers[2:end]
        z, inner = if layer isa MathOptAI.Affine
            MathOptAI.add_predictor(model, layer, [z; x])
        else
            MathOptAI.add_predictor(model, layer, z)
        end
        push!(formulation.layers, inner)
    end
    return z, formulation
end

With that, we are now ready to embed these networks into JuMP.

Embed ICNN into JuMP

Let us build a small ICNN first.

predictor = InputConvexChain(
    InputConvex((8, 8) => 2, Flux.relu),
    InputConvex((2, 8) => 1, Flux.relu),
)

InputConvexChain(
	InputConvex((8, 2) => 8, relu)
	InputConvex((2, 1) => 8, relu)
)

We can now embed predictor into a JuMP model. We choose to embed the Flux.relu using ReLUSOS1:

model = Model()
@variable(model, x[1:8])
config = Dict(Flux.relu => MathOptAI.ReLUSOS1)
z, formulation = MathOptAI.add_predictor(model, predictor, x; config);

1-element Vector{JuMP.VariableRef}:
 moai_ReLU[1]

formulation

Affine(A, b) [input: 8, output: 2]
├ variables [2]
│ ├ moai_Affine[1]
│ └ moai_Affine[2]
└ constraints [2]
  ├ 0.18672125041484833 x[1] + 0.5321521162986755 x[2] - 0.02422470599412918 x[3] + 0.4732942283153534 x[4] - 0.014558093622326851 x[5] + 0.6768441200256348 x[6] + 0.1583644598722458 x[7] + 0.28693699836730957 x[8] - moai_Affine[1] = -0.014636251144111156
  └ 0.46623796224594116 x[1] - 0.08479519933462143 x[2] + 0.7530664205551147 x[3] + 0.14653414487838745 x[4] - 0.6621298789978027 x[5] + 0.43143174052238464 x[6] - 0.7085301876068115 x[7] - 0.23360468447208405 x[8] - moai_Affine[2] = -0.8108593225479126
MathOptAI.ReLUSOS1()
├ variables [4]
│ ├ moai_ReLU[1]
│ ├ moai_ReLU[2]
│ ├ moai_z[1]
│ └ moai_z[2]
└ constraints [8]
  ├ moai_ReLU[1] ≥ 0
  ├ moai_z[1] ≥ 0
  ├ moai_Affine[1] - moai_ReLU[1] + moai_z[1] = 0
  ├ [moai_ReLU[1], moai_z[1]] ∈ MathOptInterface.SOS1{Float64}([1.0, 2.0])
  ├ moai_ReLU[2] ≥ 0
  ├ moai_z[2] ≥ 0
  ├ moai_Affine[2] - moai_ReLU[2] + moai_z[2] = 0
  └ [moai_ReLU[2], moai_z[2]] ∈ MathOptInterface.SOS1{Float64}([1.0, 2.0])
Affine(A, b) [input: 10, output: 1]
├ variables [1]
│ └ moai_Affine[1]
└ constraints [1]
  └ -0.2313876897096634 x[1] - 0.4751546084880829 x[2] + 0.6339517831802368 x[3] - 0.7303792238235474 x[4] - 0.4598158597946167 x[5] - 0.747797966003418 x[6] - 0.5654113292694092 x[7] - 0.5846427083015442 x[8] + 1.1519972085952759 moai_ReLU[1] + 1.3564424514770508 moai_ReLU[2] - moai_Affine[1] = 0.13882945477962494
MathOptAI.ReLUSOS1()
├ variables [2]
│ ├ moai_ReLU[1]
│ └ moai_z[1]
└ constraints [4]
  ├ moai_ReLU[1] ≥ 0
  ├ moai_z[1] ≥ 0
  ├ moai_Affine[1] - moai_ReLU[1] + moai_z[1] = 0
  └ [moai_ReLU[1], moai_z[1]] ∈ MathOptInterface.SOS1{Float64}([1.0, 2.0])

Epigraph formulations

The nice thing about ICNNs is that we can formulate their epigraph and avoid adding binary variables to the model. For that, we can use ReLUEpigraph.

Let's first train a model to predict the relationship $y = x^2$. (Note that this is a very basic training loop.)

Random.seed!(1234)
chain = InputConvexChain(
    InputConvex((1, 1) => 10, Flux.relu),
    InputConvex((10, 1) => 1, Flux.relu),
)

begin
    X = -2.0f0:0.1f0:2.0f0
    optimizer_state = Flux.setup(Flux.Adam(5e-2), chain)
    for epoch in 1:1000
        _, gradient = Flux.withgradient(chain) do model
            return sum((only(model([x])) - x^2)^2 for x in X)
        end
        Flux.update!(optimizer_state, chain, only(gradient))
    end
end

Now we can embed the trained network into a JuMP model:

model = Model(HiGHS.Optimizer)
set_silent(model)
@variable(model, x[1:1])
config = Dict(Flux.relu => MathOptAI.ReLUEpigraph)
y, _ = MathOptAI.add_predictor(model, chain, x; config)
@objective(model, Min, only(y))
model

A JuMP Model
├ solver: HiGHS
├ objective_sense: MIN_SENSE
│ └ objective_function_type: JuMP.VariableRef
├ num_variables: 23
├ num_constraints: 33
│ ├ JuMP.AffExpr in MOI.EqualTo{Float64}: 11
│ ├ JuMP.AffExpr in MOI.GreaterThan{Float64}: 11
│ └ JuMP.VariableRef in MOI.GreaterThan{Float64}: 11
└ Names registered in the model
  └ :x

Because we used the ReLUEpigraph predictor, there are no binary or integer variables in our model.

Moreover, we can show that the objective value y is convex with respect to x:

x_value, y_value = -2:0.1:2, Float64[]
for xi in x_value
    fix(x[1], xi)
    optimize!(model)
    # To prove we are solving an LP and not a MIP, require dual solutions.
    assert_is_solved_and_feasible(model; dual = true)
    push!(y_value, objective_value(model))
end
Plots.plot(x_value, y_value; xlabel = "x", ylabel = "y", label = "Trained")
Plots.plot!(x_value, x_value .^ 2; label = "Target", linestyle = :dash)

Conic Formulation

We can also use SoftPlusConicEpigraph in the activation functions. The resulting conic formulation can be solved using SCS or any other conic solver.

Random.seed!(1234)
chain = InputConvexChain(
    InputConvex((1, 1) => 10, Flux.softplus),
    InputConvex((10, 1) => 1, Flux.softplus),
)

begin
    X = -2.0f0:0.1f0:2.0f0
    optimizer_state = Flux.setup(Flux.Adam(5e-2), chain)
    for epoch in 1:1000
        _, gradient = Flux.withgradient(chain) do model
            return sum((only(model([x])) - x^2)^2 for x in X)
        end
        Flux.update!(optimizer_state, chain, only(gradient))
    end
end

Next, we embed the neural network using SoftPlusConicEpigraph.

model = Model(SCS.Optimizer)
set_silent(model)
@variable(model, x[1:1])
config = Dict(Flux.softplus => MathOptAI.SoftPlusConicEpigraph)
y, _ = MathOptAI.add_predictor(model, chain, x; config)
@objective(model, Min, only(y))
model

A JuMP Model
├ solver: SCS
├ objective_sense: MIN_SENSE
│ └ objective_function_type: JuMP.VariableRef
├ num_variables: 45
├ num_constraints: 44
│ ├ JuMP.AffExpr in MOI.EqualTo{Float64}: 11
│ ├ JuMP.AffExpr in MOI.LessThan{Float64}: 11
│ └ Vector{JuMP.AffExpr} in MOI.ExponentialCone: 22
└ Names registered in the model
  └ :x

Let's draw the same plot to see the differences in fit with softplus.

x_value, y_value = -2:0.1:2, Float64[]
for xi in x_value
    fix(x[1], xi)
    optimize!(model)
    # To prove we are solving an LP and not a MIP, require dual solutions.
    assert_is_solved_and_feasible(model; dual = true)
    push!(y_value, objective_value(model))
end
Plots.plot(x_value, y_value; xlabel = "x", ylabel = "y", label = "Trained")
Plots.plot!(x_value, x_value .^ 2; label = "Target", linestyle = :dash)

Nonlinear Formulation

We can also use SoftPlusEpigraph in the activation functions. The resulting global nonlinear formulation can be solved using Ipopt or any other nonlinear solver.

model = Model(Ipopt.Optimizer)
set_silent(model)
@variable(model, x[1:1])
config = Dict(Flux.softplus => MathOptAI.SoftPlusEpigraph)
y, _ = MathOptAI.add_predictor(model, chain, x; config)
@objective(model, Min, only(y))
model

A JuMP Model
├ solver: Ipopt
├ objective_sense: MIN_SENSE
│ └ objective_function_type: JuMP.VariableRef
├ num_variables: 23
├ num_constraints: 33
│ ├ JuMP.NonlinearExpr in MOI.GreaterThan{Float64}: 11
│ ├ JuMP.AffExpr in MOI.EqualTo{Float64}: 11
│ └ JuMP.VariableRef in MOI.GreaterThan{Float64}: 11
└ Names registered in the model
  └ :x

Let's draw the same plot to see the differences in fit with softplus.

x_value, y_value = -2:0.1:2, Float64[]
for xi in x_value
    fix(x[1], xi)
    optimize!(model)
    # To prove we are solving an LP and not a MIP, require dual solutions.
    assert_is_solved_and_feasible(model; dual = true)
    push!(y_value, objective_value(model))
end
Plots.plot(x_value, y_value; xlabel = "x", ylabel = "y", label = "Trained")
Plots.plot!(x_value, x_value .^ 2; label = "Target", linestyle = :dash)

This page was generated using Literate.jl.