Question about constrained problem solving.

[DoTulip]

Hello, everyone!
I am troubled by a problem and want to ask for your help. I'm doing some research with neural networks, and the objects I'm studying have some constraints (both equality and inequality constraints). Optimizers like Adam only work for unconstrained problems, what can I do to elegantly handle these constraints (I tried the penalty function method, but it doesn't seem to work very well)?

[bantin]

This doesn't seem like a jax-specific question, and it's hard to give good advice without knowing more about your problem. In general, applying constraints to the outputs of a neural net is an active area of research. Your best bet is probably to use an idea called differentiable optimization, in which the last layer of your neural net is an optimization problem with constraints. Implicit differentiation can allow you to differentiate through the solver without unrolling the loop (which is very memory expensive). JAXopt might be a good place to start.

Depending on the structure of your constraints, you might simply be able to project onto the feasible set, and differentiate through this projection step.

[DoTulip]

Thank you for your reply! It helps me a lot!