如何用PyTorch快速找到高维线性方程的N个0≤Xn≤1可行解？

核心思路 #

原随机全量生成的方法在高维空间中命中合法解的概率极低，效率极差。我们可以利用高维空间的自由度：对于n维问题，只要固定n-1个变量，剩下的1个变量可以通过方程直接求解，只需验证该变量是否在[0,1]区间内即可。这种方法的命中率远高于随机碰运气，维度越高效率提升越明显。

前提验证 #

首先需要确认问题有解：

• 最小可能的加权和为0（所有X取0）
• 最大可能的加权和为weights.sum()（所有X取1）
  如果目标值d不在[0, weights.sum()]区间内，不存在合法解。

基础实现代码 #

import torch

def generate_valid_solutions(weights, d, num_solutions=30):
    n = weights.shape[0]
    if n <= 2:
        raise ValueError("维度n必须大于2")
    
    min_d = 0.0
    max_d = weights.sum().item()
    if not (min_d <= d <= max_d):
        return torch.tensor([])  # 返回空张量表示无解
    
    # 选择第一个非零权重对应的变量作为待求解的变量，避免除以0
    non_zero_idx = (weights != 0).nonzero(as_tuple=True)[0][0]
    weights_rest = torch.cat([weights[:non_zero_idx], weights[non_zero_idx+1:]])
    
    solutions = []
    while len(solutions) < num_solutions:
        # 随机生成n-1个变量，范围[0,1]
        X_rest = torch.rand(n-1)
        # 计算已生成变量的加权和
        sum_rest = (weights_rest * X_rest).sum()
        # 求解剩余变量
        X_target = (d - sum_rest) / weights[non_zero_idx]
        
        # 验证剩余变量是否在合法区间
        if 0.0 <= X_target <= 1.0:
            # 拼接成完整的解张量
            X = torch.zeros(n)
            X[:non_zero_idx] = X_rest[:non_zero_idx]
            X[non_zero_idx] = X_target
            X[non_zero_idx+1:] = X_rest[non_zero_idx:]
            solutions.append(X)
    
    return torch.stack(solutions)

# 示例使用
weights = torch.tensor([1.0, 2.0, 3.0, 4.0])
target_d = 5.0
solutions = generate_valid_solutions(weights, target_d, 30)

# 验证解的合法性
for sol in solutions:
    assert torch.allclose((weights * sol).sum(), torch.tensor(target_d)), "解不满足约束条件"
print(f"成功生成{len(solutions)}个合法解")

批量优化版本 #

为了进一步提升效率，可以批量生成n-1个变量，一次性筛选出合法解，减少循环次数：

def generate_valid_solutions_batch(weights, d, num_solutions=30, batch_size=100):
    n = weights.shape[0]
    if n <= 2:
        raise ValueError("维度n必须大于2")
    
    min_d = 0.0
    max_d = weights.sum().item()
    if not (min_d <= d <= max_d):
        return torch.tensor([])
    
    non_zero_idx = (weights != 0).nonzero(as_tuple=True)[0][0]
    weights_rest = torch.cat([weights[:non_zero_idx], weights[non_zero_idx+1:]])
    
    solutions = []
    while len(solutions) < num_solutions:
        # 批量生成n-1维变量
        X_rest_batch = torch.rand(batch_size, n-1)
        # 批量计算加权和
        sum_rest_batch = (weights_rest * X_rest_batch).sum(dim=1)
        # 批量求解剩余变量
        X_target_batch = (d - sum_rest_batch) / weights[non_zero_idx]
        
        # 筛选合法解的掩码
        valid_mask = (X_target_batch >= 0.0) & (X_target_batch <= 1.0)
        # 提取合法的部分变量和目标变量
        valid_X_rest = X_rest_batch[valid_mask]
        valid_X_target = X_target_batch[valid_mask]
        
        # 批量拼接完整解
        for rest, target in zip(valid_X_rest, valid_X_target):
            X = torch.zeros(n)
            X[:non_zero_idx] = rest[:non_zero_idx]
            X[non_zero_idx] = target
            X[non_zero_idx+1:] = rest[non_zero_idx:]
            solutions.append(X)
    
    # 返回指定数量的解
    return torch.stack(solutions[:num_solutions])

关键注意事项 #

• 必须处理权重为0的情况：如果某权重为0，该变量的取值不影响加权和，但仍需满足0≤X≤1，因此不能选该变量作为待求解的变量，代码中通过选择第一个非零权重对应的变量避免除以0。
• 当d恰好等于0或weights.sum()时，唯一解是所有X取0或1，此时代码会自动生成全0或全1的张量。

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