如何使用 pytorch 解决严格约束的优化问题?

数据挖掘 火炬 优化
2022-02-23 13:37:31

我正在尝试使用 pytorch 解决以下问题:给定一个六面骰子,其平均滚动已知为 4.5,面的最大熵分布是多少?

(注意:我知道一堆用于解决此类问题的非 pytorch 技术——我的目标实际上是更好地理解如何使用 pytorch 解决一般的约束优化问题。在现实生活中,我正在努力工作涉及在 pytorch 中实现的神经模型的约束优化问题,我希望如果我能解决这个问题,那么它将有助于解决更难的问题。)

原则上应该可以通过寻找拉格朗日的临界点来处理这个问题:

L(p)=ipilogpi+λ(ipi1)+μ(iipi4.5)

这是我尝试使用 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()

这导致概率分布[0.1759, 0.0827, 0.0457, 0.1483, 0.2648, 0.2583]不正确——真正的答案是[0.05435, 0.07877, 0.1142, 0.1654, 0.2398, 0.3475](另外,如果我设置了,mean_constraint=3.5那么我不会得到均匀分布,所以这是一个不好的迹象。)

关于如何完成这项工作的任何想法?

1个回答

我是 Cooper 的主要贡献者,Cooper是一个专注于 Pytorch 的约束优化的库。正如您在示例中所做的那样,该库采用了约束优化问题的拉格朗日公式。

事实上,我们已经使用 Cooper“方法”作为我们自述文件中的入门片段 - :) 谢谢!我们的教程之一包含对此问题的完全可运行的答案。

虽然下面的答案侧重于离散模具示例,但 Cooper 在设计时考虑了“现实世界”神经网络问题。我们真的鼓励您查看它,如果您发现它有用/希望看到一些特定功能,请与我们联系!


为了让这个答案自成一体,这里有一种使用 Cooper 来解决这个问题的方法。(您可以使用安装 Cooper 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)

这个例子使用了一个花哨的 SGD 版本(这些是ExtraSGD 优化器),但是您可以使用 Cooper 中的几乎任何其他 Pytorch 优化器来代替。我在这里使用了 ExtraSGD,因为它有助于更​​好地控制参数振荡(参见教程中的图表)。