在 Stan/rstan 中优化高斯过程

Optimizing Gaussian Process in Stan/rstan

我最近遇到了高斯过程模型,碰巧认为它们可能是我实验室一直在研究的一个问题的解决方案。我有一个关于 Cross Validated 的开放式相关问题,但我想将我的 modeling/math 问题与我的编程问题分开。因此,这一秒,相关post。如果更多地了解我的问题的背景会有所帮助,尽管这里是 link 我的开放 CV question

这是我的 stan 代码,对应于我的 CV post 中显示的更新协方差函数:

functions{
    //covariance function for main portion of the model
    matrix main_GP(
        int Nx,
        vector x,
        int Ny,
        vector y, 
        real alpha1,
        real alpha2,
        real alpha3,
        real rho1,
        real rho2,
        real rho3,
        real rho4,
        real rho5,
        real HR_f,
        real R_f){
                    matrix[Nx, Ny] K1;
                    matrix[Nx, Ny] K2;
                    matrix[Nx, Ny] K3;
                    matrix[Nx, Ny] Sigma;

                    //specifying random Gaussian process that governs covariance matrix
                    for(i in 1:Nx){
                        for (j in 1:Ny){
                            K1[i,j] = alpha1*exp(-square(x[i]-y[j])/2/square(rho1));
                        }
                    }

                    //specifying random Gaussian process incorporates heart rate
                    for(i in 1:Nx){
                        for(j in 1:Ny){
                            K2[i, j] = alpha2*exp(-2*square(sin(pi()*fabs(x[i]-y[j])*HR_f))/square(rho2))*
                            exp(-square(x[i]-y[j])/2/square(rho3));
                        }
                    }

                    //specifying random Gaussian process incorporates heart rate as a function of respiration
                    for(i in 1:Nx){
                        for(j in 1:Ny){
                            K3[i, j] = alpha3*exp(-2*square(sin(pi()*fabs(x[i]-y[j])*HR_f))/square(rho4))*
                            exp(-2*square(sin(pi()*fabs(x[i]-y[j])*R_f))/square(rho5));
                        }
                    }

                    Sigma = K1+K2+K3;
                    return Sigma;
                }
    //function for posterior calculations
    vector post_pred_rng(
        real a1,
        real a2,
        real a3,
        real r1, 
        real r2,
        real r3,
        real r4,
        real r5,
        real HR,
        real R,
        real sn,
        int No,
        vector xo,
        int Np, 
        vector xp,
        vector yobs){
                matrix[No,No] Ko;
                matrix[Np,Np] Kp;
                matrix[No,Np] Kop;
                matrix[Np,No] Ko_inv_t;
                vector[Np] mu_p;
                matrix[Np,Np] Tau;
                matrix[Np,Np] L2;
                vector[Np] yp;

    //--------------------------------------------------------------------
    //Kernel Multiple GPs for observed data
    Ko = main_GP(No, xo, No, xo, a1, a2, a3, r1, r2, r3, r4, r5, HR, R);
    Ko = Ko + diag_matrix(rep_vector(1, No))*sn;

    //--------------------------------------------------------------------
    //kernel for predicted data
    Kp = main_GP(Np, xp, Np, xp, a1, a2, a3, r1, r2, r3, r4, r5, HR, R);
    Kp = Kp + diag_matrix(rep_vector(1, Np))*sn;

    //--------------------------------------------------------------------
    //kernel for observed and predicted cross 
    Kop = main_GP(No, xo, Np, xp, a1, a2, a3, r1, r2, r3, r4, r5, HR, R);

    //--------------------------------------------------------------------
    //Algorithm 2.1 of Rassmussen and Williams... 
    Ko_inv_t = Kop'/Ko;
    mu_p = Ko_inv_t*yobs;
    Tau=Kp-Ko_inv_t*Kop;
    L2 = cholesky_decompose(Tau);
    yp = mu_p + L2*rep_vector(normal_rng(0,1), Np);
    return yp;
    }
}

data { 
    int<lower=1> N1;
    int<lower=1> N2;
    vector[N1] X; 
    vector[N1] Y;
    vector[N2] Xp;
    real<lower=0> mu_HR;
    real<lower=0> mu_R;
}

transformed data { 
    vector[N1] mu;
    for(n in 1:N1) mu[n] = 0;
}

parameters {
    real loga1;
    real loga2;
    real loga3;
    real logr1;
    real logr2;
    real logr3;
    real logr4;
    real logr5;
    real<lower=.5, upper=3.5> HR;
    real<lower=.05, upper=.75> R;
    real logsigma_sq;
}

transformed parameters {
    real a1 = exp(loga1);
    real a2 = exp(loga2);
    real a3 = exp(loga3);
    real r1 = exp(logr1);
    real r2 = exp(logr2);
    real r3 = exp(logr3);
    real r4 = exp(logr4);
    real r5 = exp(logr5);
    real sigma_sq = exp(logsigma_sq);
}

model{ 
    matrix[N1,N1] Sigma;
    matrix[N1,N1] L_S;

    //using GP function from above 
    Sigma = main_GP(N1, X, N1, X, a1, a2, a3, r1, r2, r3, r4, r5, HR, R);
    Sigma = Sigma + diag_matrix(rep_vector(1, N1))*sigma_sq;

    L_S = cholesky_decompose(Sigma);
    Y ~ multi_normal_cholesky(mu, L_S);

    //priors for parameters
    loga1 ~ student_t(3,0,1);
    loga2 ~ student_t(3,0,1);
    loga3 ~ student_t(3,0,1);
    logr1 ~ student_t(3,0,1);
    logr2 ~ student_t(3,0,1);
    logr3 ~ student_t(3,0,1);
    logr4 ~ student_t(3,0,1);
    logr5 ~ student_t(3,0,1);
    logsigma_sq ~ student_t(3,0,1);
    HR ~ normal(mu_HR,.25);
    R ~ normal(mu_R, .03);
}

generated quantities {
    vector[N2] Ypred;
    Ypred = post_pred_rng(a1, a2, a3, r1, r2, r3, r4, r5, HR, R, sigma_sq, N1, X, N2, Xp, Y);
}

我修改了内核中包含的参数的先验,一些参数化速度更快(在某些情况下快两倍,但即使在这些情况下仍然会产生相对较慢的链)。

我正在尝试预测 3.5 秒的数据值(以 10 Hz 采样 - 所以 35 个数据点),使用来自受污染部分前后 15 秒的数据(以 3.33 Hz 采样所以总共 100 个数据)点)。

R中指定的模型如下:

 fit.pred2 <- stan(file = 'Fast_GP6_all.stan',
                 data = dat, 
                 warmup = 1000,
                 iter = 1500,
                 refresh=5,
                 chains = 3,
                 pars= pars.to.monitor
                 )

老实说,我不知道我是否需要那么多预热迭代。我认为估计缓慢的部分原因是先验信息相当少(心率和呼吸 HRR 除外,因为它们在健康成年人的静息状态下具有相当众所周知的范围)。

欢迎提出任何建议来加快我程序的 运行 时间。

谢谢。

您可以获取 Stan 数学库的开发分支,该分支最近更新了 multi_normal_cholesky 版本,该版本在内部使用解析梯度而不是自动微分。为此,您可以在 R 中执行 source("https://raw.githubusercontent.com/stan-dev/rstan/develop/install_StanHeaders.R") 但是你需要在你的 ~/.R/Makevars 文件中包含 CXXFLAGS+=-std=c++11 并且之后可能需要重新安装 rstan 包。

multi_normal_cholesky 和你的 main_GP 都是 O(N^3),所以你的 Stan 程序永远不会特别快,但是这两个函数的增量优化将使最大差异。

除此之外,还有一些小东西,比如 Sigma = Sigma + diag_matrix(rep_vector(1, N1))*sigma_sq; 应该改成 for (n in 1:N1) Sigma[n,n] += sigma_sq; 原因是将 sigma_sq 乘以对角矩阵会将 N1 平方节点放入 autodiff 树,将其添加到 Sigma 也是如此,这会进行大量内存分配和释放。沿对角线的显式循环仅将 N1 节点放入自动差异树,或者如果我们巧妙地使用 += 运算符,它可能只是更新现有树。在你的 post_pred_rng 函数中有同样的处理,尽管这不是那么重要,因为生成的数量块是用双精度而不是反向模式自动差异的自定义 Stan 类型来计算的。

我认为/希望 vector[N2] Ypred = post_pred_rng(...); vector[N2] Ypred; // preallocates Ypred with NaNs Ypred = post_pred_rng(...); 通过避免预分配步骤;无论如何,读起来更好。

最后,虽然它不影响速度,但您没有义务以对数形式指定参数,然后在转换后的参数块中反对它们。您可以只用 <lower=0> 声明事物,这将导致相同的结果,尽管您将在积极约束的事物而不是不受约束的事物上指定您的先验。在大多数情况下,这更直观。那些具有 3 个自由度的 Student t 先验是非常重尾的,这可能会导致 Stan 至少在热身期间采取大量的蛙跳步数(默认情况下达到 10 步的限制)。蛙跳步数 (s) 是运行时间的主要贡献者,因为每次迭代都需要 2^s - 1 函数/梯度评估。