基于GPyTorch的回归任务与梯度符号约束建模技术求助

针对你需要的4维输入回归任务（约束3个输入梯度为正、1个为负），以下提供两种可行的实现方案，替代你当前效果不佳的手动梯度损失方法：

方案一：GPyTorch多任务多似然实现 #

核心逻辑 #

构建多任务高斯过程，同时建模原始回归输出和4个输入维度的梯度，用高斯似然处理回归任务，伯努利Probit似然处理梯度符号约束，利用GP的联合后验自动平衡任务权重。

1. 定义多任务核函数 #

使用DerivativeKernel生成梯度对应的核，结合多任务核组合函数值与所有梯度输出：

import torch
import gpytorch
from gpytorch.kernels import RBFKernel, ScaleKernel, DerivativeKernel, MultitaskKernel
from gpytorch.means import ConstantMean
from gpytorch.models import ExactGP
from gpytorch.likelihoods import GaussianLikelihood, BernoulliLikelihood, MultitaskLikelihood
from gpytorch.mlls import SumMarginalLogLikelihood

# 多任务模型：输出维度=1（函数值）+4（4个输入维度的梯度）=5
class MultitaskGradGP(ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super().__init__(train_x, train_y, likelihood)
        self.mean_module = ConstantMean(batch_shape=torch.Size([5]))
        # 基础ARD核适配4维输入
        base_kernel = RBFKernel(ard_num_dims=4)
        # 函数值对应的核
        func_kernel = ScaleKernel(base_kernel)
        # 4个输入维度的导数核
        grad_kernels = [ScaleKernel(DerivativeKernel(base_kernel, dim=i)) for i in range(4)]
        
        # 组合多任务核
        self.covar_module = MultitaskKernel(
            func_kernel,
            task_covar_module=gpytorch.kernels.LinearKernel(num_tasks=5),
            num_tasks=5
        )
    
    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)

2. 构造多似然与训练数据 #

• 高斯似然处理原始回归目标
• 伯努利Probit似然处理梯度符号：正梯度标签设为1，负梯度标签设为0

# 准备训练数据：
# train_x: (N,4) 4维输入样本
# train_y_func: (N,1) 回归任务目标值
# train_y_grad_sign: (N,4) 梯度符号标签，1=正，0=负

# 组合多任务目标
train_y = torch.cat([train_y_func, train_y_grad_sign], dim=-1)

# 定义双似然组合
likelihood_func = GaussianLikelihood()
likelihood_grad = BernoulliLikelihood(link=gpytorch.likelihoods.links.ProbitLink())
likelihood = MultitaskLikelihood(
    [likelihood_func] + [likelihood_grad]*4,
    num_tasks=5
)

# 初始化模型与损失
model = MultitaskGradGP(train_x, train_y, likelihood)
mll = SumMarginalLogLikelihood(likelihood, model)

3. 训练循环 #

model.train()
likelihood.train()
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

for epoch in range(1000):
    optimizer.zero_grad()
    output = model(train_x)
    loss = -mll(output, train_y)
    loss.backward()
    optimizer.step()
    if epoch % 100 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

方案二：复现Riihimäki论文的高维适配方案 #

论文用期望传播（EP）处理混合似然，GPy的EP仅支持一维，可改用Pyro的变分EP或SVI实现高维适配：

核心逻辑 #

将梯度符号约束转化为隐变量模型，用变分近似拟合联合后验，同时建模函数值与梯度的依赖关系：

import pyro
import pyro.distributions as dist
from pyro.infer import SVI, Trace_ELBO
from pyro.optim import Adam
import gpytorch

def model(x, y_func, y_grad_sign):
    # 核参数先验
    lengthscale = pyro.sample("lengthscale", dist.LogNormal(0, 1).expand([4]))
    variance = pyro.sample("variance", dist.LogNormal(0, 1))
    
    # 构建GP先验
    kernel = gpytorch.kernels.RBFKernel(ard_num_dims=4).to(x)
    kernel.lengthscale = lengthscale
    kernel.variance = variance
    gp_cov = kernel(x).evaluate() + 1e-6 * torch.eye(len(x))
    f = pyro.sample("f", dist.MultivariateNormal(torch.zeros(len(x)), gp_cov))
    
    # 回归似然
    with pyro.plate("data", len(x)):
        pyro.sample("y_func", dist.Normal(f, 0.05), obs=y_func.squeeze())
        
        # 自动微分计算梯度
        f_grad = torch.autograd.grad(f.sum(), x, create_graph=True)[0]
        
        # 梯度符号的伯努利Probit似然
        for dim in range(4):
            prob = dist.Normal(0,1).cdf(f_grad[:,dim])
            pyro.sample(f"y_grad_{dim}", dist.Bernoulli(prob), obs=y_grad_sign[:,dim])

def guide(x, y_func, y_grad_sign):
    # 变分近似GP后验
    f_loc = pyro.param("f_loc", torch.zeros(len(x)))
    f_scale_tril = pyro.param("f_scale_tril", torch.eye(len(x)) + 1e-6)
    pyro.sample("f", dist.MultivariateNormal(f_loc, scale_tril=f_scale_tril))
    
    # 核参数变分近似
    lengthscale_loc = pyro.param("lengthscale_loc", torch.zeros(4))
    lengthscale_scale = pyro.param("lengthscale_scale", torch.ones(4))
    pyro.sample("lengthscale", dist.LogNormal(lengthscale_loc, lengthscale_scale))
    
    variance_loc = pyro.param("variance_loc", torch.tensor(0.))
    variance_scale = pyro.param("variance_scale", torch.tensor(1.))
    pyro.sample("variance", dist.LogNormal(variance_loc, variance_scale))

# 训练
optimizer = Adam({"lr": 0.01})
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

for epoch in range(1000):
    loss = svi.step(train_x, train_y_func, train_y_grad_sign)
    if epoch % 100 == 0:
        print(f"Epoch {epoch}, Loss: {loss/len(train_x):.4f}")

现有方案效果不佳的原因 #

你手动计算梯度并加权Probit损失的方法存在两个核心问题：

1. 未利用GP的联合后验信息，梯度的方差估计与函数值后验完全分离
2. 手动设置的损失权重（100倍）极易导致训练不稳定，多似然框架会自动平衡不同任务的损失贡献

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