3
$\begingroup$

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()
    optimizer.zero_grad()
    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?

$\endgroup$

1 Answer 1

4
$\begingroup$

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):
    coop.zero_grad()
    lagrangian = formulation.composite_objective(cmp.closure, probs)
    formulation.custom_backward(lagrangian)
    coop.step(cmp.closure, probs)

This example uses a fancy version of SGD (these are the ExtraSGD optimizers), but you could use almost any other Pytorch optimizer in Cooper instead. I used ExtraSGD here since it helped control the parameter oscillations better (see plots in tutorial).

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.