RcppParallel适配RcppArmadillo的随机种子问题求助

问题背景 #

我有一段可正常运行的RcppArmadillo串行代码，尝试通过RcppParallel适配为并行版本后，速度有所提升但计算结果不正确，推测是并行计算时随机种子使用不当导致的。

原代码如下：

// Worker 
struct UpdateOmegaWorker : public Worker {
const arma::mat X;
const arma::mat Omega_old;
const arma::vec delta;
const double W_slab;
const double W_spike;
const double seed_for_start;

arma::mat& Omega_ret;

// Costruttore

UpdateOmegaWorker(const arma::mat X, const arma::mat Omega_old, const arma::vec delta,
const double W_slab, const double W_spike, arma::mat& Omega_ret, const double seed_for_start): X(X), Omega_old(Omega_old), delta(delta), W_slab(W_slab), W_spike(W_spike), Omega_ret(Omega_ret), seed_for_start(seed_for_start) {}

// Operator
void operator()(std::size_t begin, std::size_t end) {
int n = X.n_rows;
int p = X.n_cols;

    arma::mat Xs_it = X;
    arma::colvec x_it(n);
    arma::colvec phi(n);
    arma::colvec pg_aux(n);
    arma::colvec kappa_it(n);
    arma::colvec zeta_it(n);
    arma::colvec xs_tilde(n);
    
    arma::mat node_V(p, p);
    arma::colvec node_m(p);
    arma::mat Winv = arma::diagmat(1 / (W_slab * delta + W_spike * (1 - delta)));
    

    
    for (std::size_t s = begin; s < end; s++) {
      
      // HERE IS PROBABILY THE ERROR ?! 
      arma::arma_rng::set_seed( seed_for_start ) ;
      
      //cores Code (copy Paste of a working code)
      // Parametrizzazione
      Xs_it.col(s).fill(1.0);
      xs_tilde = X.col(s);
      
      // Pòlya Gamma Augmentation
      zeta_it = Xs_it * Omega_old.col(s);
      pg_aux = cpp_rpg_z_vec(n, 1, zeta_it, 200);
      
      // Ising-Regression
      Xs_it.col(s) = xs_tilde;
      kappa_it = xs_tilde - 0.5;
      
      node_V = arma::inv(Xs_it.t() * (X.each_col() % pg_aux) + Winv);
      node_m = node_V * Xs_it.t() * kappa_it;
      
      Omega_ret.col(s) = cpp_mvrnormArma1(node_m, node_V);
    }
  }
};

// [[Rcpp::export]]
arma::mat update_Omega_parallel(const arma::mat X, const arma::mat Omega_old,
const arma::vec delta, const double W_slab, const double W_spike, const double seed_for_start) {
int p = X.n_cols;
arma::mat Omega_ret(p, p, arma::fill::zeros);
UpdateOmegaWorker worker(X, Omega_old, delta, W_slab, W_spike, Omega_ret, seed_for_start);
parallelFor(0, p, worker);

return Omega_ret;
}

问题分析 #

原代码的核心问题在于：

• 所有并行线程在每次迭代时都将随机种子设置为同一个固定值seed_for_start，导致不同线程生成完全相同的随机序列，最终计算结果偏离预期。
• Armadillo的全局随机数生成器并非线程安全，多线程同时修改全局种子会引发竞争条件，进一步加剧结果的不确定性。

修复方案 #

采用原子计数器为每个迭代分配唯一的随机种子，确保线程安全且随机序列独立。具体修改如下：

修改后的代码 #

// Worker 
struct UpdateOmegaWorker : public Worker {
    const arma::mat& X;
    const arma::mat& Omega_old;
    const arma::vec& delta;
    const double W_slab;
    const double W_spike;
    std::atomic<std::size_t>& counter; // 共享原子计数器，用于生成唯一种子
    arma::mat& Omega_ret;

    // 构造函数
    UpdateOmegaWorker(const arma::mat& X, const arma::mat& Omega_old, const arma::vec& delta,
                      const double W_slab, const double W_spike, arma::mat& Omega_ret, std::atomic<std::size_t>& counter)
        : X(X), Omega_old(Omega_old), delta(delta), W_slab(W_slab), W_spike(W_spike), Omega_ret(Omega_ret), counter(counter) {}

    // 并行执行算子
    void operator()(std::size_t begin, std::size_t end) {
        int n = X.n_rows;
        int p = X.n_cols;

        arma::mat Xs_it = X;
        arma::colvec x_it(n);
        arma::colvec phi(n);
        arma::colvec pg_aux(n);
        arma::colvec kappa_it(n);
        arma::colvec zeta_it(n);
        arma::colvec xs_tilde(n);
        
        arma::mat node_V(p, p);
        arma::colvec node_m(p);
        arma::mat Winv = arma::diagmat(1 / (W_slab * delta + W_spike * (1 - delta)));
        
        for (std::size_t s = begin; s < end; s++) {
            // 为当前迭代生成唯一种子
            std::size_t current_seed = counter++;
            arma::arma_rng::set_seed(current_seed);
            
            // 原有核心计算代码
            Xs_it.col(s).fill(1.0);
            xs_tilde = X.col(s);
            
            // Pòlya Gamma Augmentation
            zeta_it = Xs_it * Omega_old.col(s);
            pg_aux = cpp_rpg_z_vec(n, 1, zeta_it, 200);
            
            // Ising-Regression
            Xs_it.col(s) = xs_tilde;
            kappa_it = xs_tilde - 0.5;
            
            node_V = arma::inv(Xs_it.t() * (X.each_col() % pg_aux) + Winv);
            node_m = node_V * Xs_it.t() * kappa_it;
            
            Omega_ret.col(s) = cpp_mvrnormArma1(node_m, node_V);
        }
    }
};

// [[Rcpp::export]]
arma::mat update_Omega_parallel(const arma::mat& X, const arma::mat& Omega_old,
                                const arma::vec& delta, const double W_slab, const double W_spike, const std::size_t seed_for_start) {
    int p = X.n_cols;
    arma::mat Omega_ret(p, p, arma::fill::zeros);
    // 初始化原子计数器，起始值为传入的种子
    std::atomic<std::size_t> counter(seed_for_start);
    UpdateOmegaWorker worker(X, Omega_old, delta, W_slab, W_spike, Omega_ret, counter);
    parallelFor(0, p, worker);

    return Omega_ret;
}

关键修改点 #

1. 原子计数器替代固定种子：使用std::atomic<std::size_t>作为共享计数器，确保多线程下自增操作线程安全，每个迭代获得唯一的种子值。
2. 常量参数改用引用：将X、Omega_old等大型矩阵参数改为const&引用，避免并行时的不必要内存复制，提升效率。
3. 独立种子生成：每个循环迭代使用递增的种子值，保证不同线程生成的随机序列完全独立，消除结果重复问题。

额外注意事项 #

• 如果cpp_rpg_z_vec或cpp_mvrnormArma1内部依赖全局随机状态，建议为每个线程创建独立的Armadillo随机数生成器实例（例如arma::rng::generator的局部实例），进一步规避线程安全风险。
• 确保编译时启用C++11及以上标准，std::atomic需要该版本支持。

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