如何在GPyTorch中实现与Scikit-learn等效的多目标高斯过程回归？

问题背景 #

我们需要实现函数到函数的高斯过程回归：将N个离散采样（p个点）的函数构造成N×p的输入矩阵X，逐行映射到同形状的输出矩阵Y。在Scikit-learn中，仅需单个核即可快速搭建多目标GP模型，代码如下：

from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF

kernel = RBF(1.0)
gp = GaussianProcessRegressor(kernel=kernel, n_restarts_optimizer=10)
gp.fit(X, Y)

但在GPyTorch中使用IndependentModelList共享核实例的方式，无法达到同等效果，且运行效率更低。

问题分析 #

Scikit-learn的多目标GP实现中，若使用单输出核，实际是对每个输出维度拟合独立GP但共享核参数；若要显式捕捉目标间的相关性，则需使用多输出核。而之前的GPyTorch代码用IndependentModelList时，即使共享核与噪声参数，每个模型仍是独立的单任务GP，没有联合建模输出维度间的协方差，因此无法捕捉目标间的相关性，导致效果不佳。

正确的做法是使用GPyTorch的多输出（多任务）GP模型，显式建模目标间的相关性。

等效实现代码 #

import torch
import gpytorch
import numpy as np
import matplotlib.pyplot as plt
from scipy.special import factorial

# 1. 生成数据
def schulz(x, mean, z):
    return 1/factorial(z)*((z+1)/mean)**(z+1)*x**z*np.exp(-(z+1)/mean*x)

npoints = 64  # 每个函数的采样点数量p
tmin, tmax = 0, 10 
t = np.linspace(tmin, tmax, npoints)

Nf = 200  # 训练函数数量N
zmin = 1
zmax = 10

# 生成输入函数（N×p）
X = np.array([schulz(t, np.random.uniform(0, tmax/2.5), np.random.randint(zmin, zmax)) for _ in range(Nf)])  

def transform(u, eps=1e-3):
    F = np.log(u+eps)
    return F + F**2 + F**3

# 生成输出函数（N×p）
Y = np.array([transform(_) for _ in X])  

# 测试函数
f_test = schulz(t, 2.1, zmax)
y_test_true = transform(f_test)

# 2. 数据转换为多任务GP格式
# 多任务GP需要将数据整理为：(N*p, 2)的输入，其中第二列是任务索引（0到p-1）
train_x = []
train_y = []
for task_idx in range(npoints):
    # 每个任务对应X和Y的一列
    task_x = X[:, task_idx].reshape(-1, 1)
    # 添加任务索引列
    task_x_with_idx = np.hstack([task_x, np.full_like(task_x, task_idx)])
    train_x.append(task_x_with_idx)
    train_y.append(Y[:, task_idx])

train_x = torch.tensor(np.vstack(train_x), dtype=torch.float32)
train_y = torch.tensor(np.concatenate(train_y), dtype=torch.float32)

# 测试数据转换：每个采样点对应一个任务
test_x = []
for task_idx in range(npoints):
    test_point = np.array([[f_test[task_idx], task_idx]])
    test_x.append(test_point)
test_x = torch.tensor(np.vstack(test_x), dtype=torch.float32)

# 3. 定义多任务GP模型
class MultiTaskGPModel(gpytorch.models.ExactMultiTaskGP):
    def __init__(self, train_x, train_y, num_tasks):
        # 多任务似然，支持不同任务的噪声（也可共享）
        likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)
        super().__init__(train_x, train_y, likelihood)
        # 均值函数：每个任务共享常数均值
        self.mean_module = gpytorch.means.MultitaskMean(
            gpytorch.means.ConstantMean(), num_tasks=num_tasks
        )
        # 核函数：RBF核（作用于输入数值） + 线性核（作用于任务索引，捕捉任务间相关性）
        self.covar_module = gpytorch.kernels.MultitaskKernel(
            gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(active_dims=[0])),
            num_tasks=num_tasks,
            rank=1,  # 任务间协方差矩阵的秩，控制相关性复杂度
            active_dims=[1]
        )

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)

# 初始化模型
model = MultiTaskGPModel(train_x, train_y, num_tasks=npoints)
likelihood = model.likelihood

# 4. 训练模型
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)
model.train()
likelihood.train()

optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
training_iterations = 100

for i in range(training_iterations):
    optimizer.zero_grad()
    output = model(train_x)
    loss = -mll(output, train_y)
    loss.backward()
    if (i+1) % 10 == 0:
        print(f'Iter {i+1}/{training_iterations} - Loss: {loss.item():.3f}')
    optimizer.step()

# 5. 预测
model.eval()
likelihood.eval()

with torch.no_grad(), gpytorch.settings.fast_pred_var():
    observed_pred = likelihood(model(test_x))

# 提取预测结果
pred_mean = observed_pred.mean.numpy()
lower, upper = observed_pred.confidence_region()
lower = lower.numpy()
upper = upper.numpy()

# 可视化
plt.figure(figsize=(10,6))
plt.plot(t, y_test_true, label='真实输出', color='black')
plt.plot(t, pred_mean, label='预测均值', color='red')
plt.fill_between(t, lower, upper, alpha=0.3, label='95%置信区间')
plt.legend()
plt.xlabel('t')
plt.ylabel('输出值')
plt.title('多任务GP函数预测')
plt.show()

关键说明 #

1. 数据格式转换：多任务GP需要将每个输出维度视为一个“任务”，因此输入需添加任务索引列，将N×p的X和Y转换为N*p×2和N*p的张量。
2. 多任务核设计：使用MultitaskKernel组合RBF核（建模输入数值的相关性）和任务核（建模不同输出维度间的相关性），其中rank参数控制任务间协方差矩阵的复杂度。
3. 模型选择：ExactMultiTaskGP是GPyTorch专门为多输出场景设计的精确GP模型，自动处理多任务的联合似然计算，既可以匹配Scikit-learn共享核参数的逻辑，也能显式捕捉目标间的相关性，效果更优。

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