# How do you solve strictly constrained optimization problems with pytorch?

I am trying to solve the following problem using pytorch: given a six sided die whose average roll is known to be 4.5, what is the maximum entropy distribution for the faces?

(Note: I know a bunch of non-pytorch techniques for solving problems of this sort - my goal here is really to be better understand how to solve constrained optimization problems in general with pytorch. In real life I'm working on a much harder constrained optimization problem involving a neural model implemented in pytorch, and I'm hoping that if I can solve this problem then it will help with the harder problem.)

In principle it should be possible to handle this by looking for critical points of the Lagrangian:

$$L(p) = -\sum_i p_i \log p_i + \lambda\left(\sum_i p_i - 1\right) + \mu\left(\sum_i i p_i - 4.5\right)$$

Here's my attempt to do this with pytorch:

class MaxEntropyDice(torch.nn.Module):
def __init__(self, num_faces=6, mean_constraint=3.5):
super().__init__()
self.num_faces = num_faces
self.mean_constraint = mean_constraint
self.p = torch.nn.Parameter(F.normalize(torch.rand(num_faces), p=1, dim=0))
self.probability_multiplier = torch.nn.Parameter(torch.rand(1))
self.mean_multiplier = torch.nn.Parameter(torch.rand(1))

def forward(self):
entropy = -torch.sum(self.p * torch.log(self.p))
probability_term = self.probability_multiplier * (torch.sum(self.p) - 1)
mean_term = self.mean_multiplier * (
torch.sum(torch.tensor(range(1, self.num_faces + 1)) * self.p) - self.mean_constraint
)
lagrangian = entropy + probability_term + mean_term
return lagrangian

model = MaxEntropyDice(num_faces=6, mean_constraint=4.5)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-6)

for i in range(2000):
loss = model()
loss.backward()
optimizer.step()


This results in the probability distribution [0.1759, 0.0827, 0.0457, 0.1483, 0.2648, 0.2583], which is not correct - the true answer is [0.05435, 0.07877, 0.1142, 0.1654, 0.2398, 0.3475]. (Also, if I set mean_constraint=3.5 then I don't get the uniform distribution, so that's a bad sign.)

Any ideas on how I can make this work?

I am the lead contributor to Cooper, a library focused on constrained optimization for Pytorch. The library employs a Lagrangian formulation of the constrained optimization problem, as you do in your example.

In fact, we have used the Cooper "approach" to your question as the getting started snippet in our README -- :) thanks! One of our tutorials contains a fully-runnable answer to this question.

While the answer below focuses on the discrete die example, Cooper is designed with "real world" neural net problems in mind. We really encourage you to check it out, and get in touch if you find it useful/would like to see some specific features!

To keep this answer self-contained, here is a way to approach this problem using Cooper. (You can install Cooper using pip install git+https://github.com/cooper-org/cooper.git)

import torch
import cooper

class MaximumEntropy(cooper.ConstrainedMinimizationProblem):
def __init__(self, mean_constraint):
self.mean_constraint = mean_constraint
super().__init__(is_constrained=True)

def closure(self, probs):
# Verify domain of definition of the functions
assert torch.all(probs >= 0)

# Negative sign removed since we want to *maximize* the entropy
neg_entropy = torch.sum(probs * torch.log(probs))

# Entries of p >= 0 (equiv. -p <= 0)
ineq_defect = -probs

# Equality constraints for proper normalization and mean constraint
mean = torch.sum(torch.tensor(range(1, len(probs) + 1)) * probs)
eq_defect = torch.stack([torch.sum(probs) - 1, mean - self.mean_constraint])

return cooper.CMPState(loss=neg_entropy, eq_defect=eq_defect, ineq_defect=ineq_defect)

# Define the problem and formulation
cmp = MaximumEntropy(mean_constraint=4.5)
formulation = cooper.LagrangianFormulation(cmp)

# Define the primal parameters and optimizer
rand_init = torch.rand(6)  # Use a 6-sided die
probs = torch.nn.Parameter(rand_init / sum(rand_init))
primal_optimizer = cooper.optim.ExtraSGD([probs], lr=3e-2, momentum=0.7)

# Define the dual optimizer. Note that this optimizer has NOT been fully instantiated
# yet. Cooper takes care of this, once it has initialized the formulation state.
dual_optimizer = cooper.optim.partial(cooper.optim.ExtraSGD, lr=9e-3, momentum=0.7)

# Wrap the formulation and both optimizers inside a ConstrainedOptimizer
coop = cooper.ConstrainedOptimizer(formulation, primal_optimizer, dual_optimizer)

# Here is the actual training loop
for iter_num in range(5000):