pytorch中概率单纯形的投影梯度下降

Projected gradient descent on probability simplex in pytorch

我有一个尺寸为 1000x70000 的矩阵 A。 我的损失函数包括 A,我想使用梯度下降找到 A 的最优值,其中约束是 A 的行保持在概率单纯形中(即每行总和为 1)。我已经初始化 A 如下所示

A=np.random.dirichlet(np.ones(70000),1000)
A=torch.tensor(A,requires_grad=True)

我的训练循环如下所示

for epoch in range(500):
    y_pred=forward(X)
    y=model(torch.mm(A.float(),X))
    l=loss(y,y_pred)
    l.backward()
    A.grad.data=-A.grad.data
    optimizer.step()
    optimizer.zero_grad()
    if epoch%2==0:
        print("Loss",l,"\n")

一个简单的方法是不直接使用 A 进行计算,而是使用 A.

的行规范化版本
# you can keep 'A' unconstrained
A = torch.rand(1000, 70000, requires_grad=True)

然后将每一行除以其总和(保持行总和始终为 1)

for epoch in range(500):
    y_pred = forward(X)
    B = A / A.sum(-1, keepdim=True) # normalize rows manually
    y = model(torch.mm(B, X))
    l = loss(y,y_pred)
    ...

所以现在,在每个步骤中,B 是约束矩阵 - 即您感兴趣的数量。但是,优化仍将进行(不受约束)A.


编辑:@Umang Gupta 在评论部分提醒我,OP 想要“概率单纯形”,这意味着会有另一个约束,即 A >= 0.

为此,您可以在每次迭代

A 上简单地应用一些适当的激活函数(例如torch.exptorch.sigmoid
A_ = torch.exp(A)
B = A_ / A_.sum(-1, keepdim=True) # normalize rows

函数的确切选择取决于需要试验的训练动力学行为。