JuMP - Nested Optimization with user defined functions

Hi Odow,

I wonder how to register an inner JuMP optimization as the nested function and use the user-defined function in constraint or objective of the outer JuMP optimization. The motivation for doing this lies in the feature of general economic models: A hacky guide to using automatic differentiation in nested optimization problems - #6 by jeffreyesun.

I am working with a maximum likelihood case where the likelihood sums across M markets

\max_{\theta, \delta}\sum_{m=1}^M l(\theta,\delta_m)

\theta is a low dimensional vector with 5 to 10 parameters. \delta are market-m specific; each market \delta_m contains about 15 parameters that correspond to available products in the market m.

Writing the whole optimization in one JuMP instance with sparse index

function mle(data)
    model = Model(Ipopt.Optimizer)
    @variable(model, θ[1:5])
    @variable(model, δ[m in market_idset, j in product_market[m]])
    @variable(model, logli[m in market_idset])
    ...
    @NLobjective(model, Min, -sum(logli[m] for m in market_idset))
end

It works well when M is small. (less than 4 minutes with M=5 markets, 300,000 total variables, 16G MacBook Pro). However, my full sample has 300~400 markets and JuMP would never start the iteration. The total number of variables reaches at a scale of ten million in the full sample case. I am considering whether a nested optimization could improve the performance.

For each market, the log likelihood l(\delta_m,\theta) is convex in \delta_m. Fixing \theta, \max_{\delta_m}l(\delta_m,\theta) takes about 10 seconds to find the local optimal solution for \delta_m. I would like to rephrase the JuMP optimization in the light of the nested structure. An incomplete/incorrect JuMP I am trying to build from your example follows:

function outer_f(data)
    outer_model = Model(Ipopt.Optimizer)
    @variable(outer_model, θ[1:5])
    @variable(outer_model, logli[m in market_idset])
    @register(outer_model,:inner_f,5,inner_JuMP_delta)
    for m in mkt_idset
        @constraint(outer_model,logli[m] == inner_f(θ)[m][1])
    end
    @NLobjective(outer_model,Min,-sum(logli[m] for m in mkt_idset))
    optimize(outer_model)
end


function inner_JuMP_delta(para...,)
    result_obj = zeros(num_markets)
    result_delta = zeros(num_markets, num_products)
    for m in market_idset
        inner_model = Model(Ipopt.Optimizer)
        @variable(inner_model, δ in product_market[m])
        ...
        @NLobjective(log(...))
        obj_result[m] = objective_value(inner_model)
        result_prod[m,:] = value.(δ)
    end
    result, result_delta
end

The main issue is that how to register a rather complicated user-definition function involving nonlinear expression, data structure (like Dictionary I am heavily using) and also returning a large array. If JuMP couldn’t work, I am considering using Optim,jl in the outer optimization. I tried it with M=5 and the derivative-free approach, but it was significantly slower than the one JuMP optimization mle().

Thanks for any thought!