Constrain sum of weights to maximum value during training

Hi all,

Hoping to get some help with a tricky optimization problem. I’m using PyTorch to minimize a function f subject to certain constraints. The constraints are that the coefficients of the function all must be >= 0, and the sum of coefficients must be <= some value W. With projected gradient descent, it’s easy enough to handle the first constraint, and indeed, it works as expected.

However, I’m having trouble coming up with a way to enforce the constraint that the sum of coefficients must be <= W. One thought was to essentially apply a large penalty whenever W - sum(coefficients) > 0. However, I’m not sure whether this is an “appropriate” way to solve this issue and what penalty function to apply. Ideally it’d be something that is 0 when x <= 0 and some exponentially increasing value when x > 0.

Without the 2nd constraint, my code is:

    opt = optim.SGD([x], lr=1)
    for e in range(max_iter):
        loss = f(x, *args)
        opt.zero_grad()
        loss.backward()
        opt.step()
        with torch.no_grad():
            x.clamp_(0, 1e10)
    return x

Any guidance or tips would be much appreciated!