权重条件数有界的感知机训练约束在PyTorch下的实现方法问询

目前存在两类成熟的标准实现方案，可根据对约束的严格程度要求选择：

• 硬约束方案：梯度更新后投影，可严格保证每一步权重的条件数低于阈值k0
  实现逻辑为每次权重更新完成后，对权重矩阵做SVD分解，调整奇异值使得最大奇异值与最小奇异值的比值不超过k0，再重构回权重矩阵。该方案的缺点是当n较大时，SVD计算会带来一定的性能开销。
• 软约束方案：损失函数加正则项，训练速度更快，适合对约束严格度要求不高的场景
  实现逻辑为在原始损失的基础上，添加条件数超过k0的惩罚项，引导优化过程自动将条件数控制在阈值以下，缺点是无法100%保证每一步都满足约束，需要在训练过程中监控条件数指标。

PyTorch 实现示例 #

硬约束实现代码 #

import torch
import torch.nn as nn
import torch.optim as optim

def project_cond_constraint(w: torch.Tensor, k0: float) -> torch.Tensor:
    # SVD分解权重矩阵
    U, singular_vals, Vh = torch.linalg.svd(w, full_matrices=False)
    current_cond = singular_vals.max() / singular_vals.min()
    # 已经满足约束直接返回原权重
    if current_cond <= k0:
        return w
    # 调整最小奇异值，保证条件数不超过k0
    min_singular_target = singular_vals.max() / k0
    adjusted_singular = torch.clamp(singular_vals, min=min_singular_target)
    # 重构权重矩阵
    return U @ torch.diag(adjusted_singular) @ Vh

# 训练流程示例
n = 16 # 权重矩阵维度
k0 = 10 # 条件数阈值
model = nn.Linear(n, n, bias=False) # 线性单层感知机
optimizer = optim.SGD(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()

for batch_x, batch_y in train_dataloader:
    optimizer.zero_grad()
    pred = model(batch_x)
    loss = loss_fn(pred, batch_y)
    loss.backward()
    optimizer.step()
    # 每次权重更新后执行投影约束
    with torch.no_grad():
        model.weight.copy_(project_cond_constraint(model.weight, k0))

软约束实现代码 #

def cond_penalty(w: torch.Tensor, k0: float, coef: float = 1e-2) -> torch.Tensor:
    singular_vals = torch.linalg.svdvals(w)
    current_cond = singular_vals.max() / singular_vals.min()
    # 超过阈值才加惩罚
    return coef * torch.relu(current_cond - k0)

# 训练时修改损失计算逻辑即可
for batch_x, batch_y in train_dataloader:
    optimizer.zero_grad()
    pred = model(batch_x)
    ori_loss = loss_fn(pred, batch_y)
    # 加条件数惩罚项
    loss = ori_loss + cond_penalty(model.weight, k0)
    loss.backward()
    optimizer.step()

内容的提问来源于stack exchange，提问作者rmcerafl