如何优化多次循环执行的NumPy矩阵点积运算？

核心思路 #

你当前的代码是重复执行 np.dot(A, W) 共n次，本质等价于计算 Aⁿ · W。直接循环的时间复杂度为O(n*M²)，当n较大时效率极低，以下是两种高效优化方式：

• 矩阵快速幂法
  利用快速幂思想，将矩阵幂的计算复杂度降至O(logn * M³)，步骤如下：

  1. 实现矩阵快速幂函数，计算A^n
  2. 用得到的A^n与W做一次点积

  示例代码：

  def matrix_pow(mat, power):
      # 初始化结果为单位矩阵
      result = np.eye(mat.shape[0])
      while power > 0:
          if power % 2 == 1:
              result = np.dot(result, mat)
          mat = np.dot(mat, mat)
          power = power // 2
      return result
  
  # 计算A^n后与W相乘
  A_n = matrix_pow(A, n)
  W = np.dot(A_n, W)
  

• 特征值分解法
  若矩阵A可对角化（多数实际场景满足），可通过特征值分解简化计算：

  1. 对A做特征值分解：A = Q @ Λ @ Q_inv，其中Λ是对角矩阵，元素为A的特征值
  2. Aⁿ = Q @ Λⁿ @ Q_inv，Λⁿ只需将对角元素取n次幂，此步骤复杂度为O(M)
  3. 最终计算W = Q @ (Λⁿ @ (Q_inv @ W))

  示例代码：

  # 执行特征值分解
  eigenvalues, Q = np.linalg.eig(A)
  Q_inv = np.linalg.inv(Q)
  
  # 分步计算中间结果
  temp = Q_inv @ W
  temp = temp * (eigenvalues ** n).reshape(-1, 1)
  
  # 得到最终结果
  W = Q @ temp
  

方案选择 #

• 当n极大时，特征值分解法效率更高，矩阵幂计算被转化为简单的元素幂运算
• 若矩阵A不可对角化，或n较小，矩阵快速幂法更稳妥，可避免分解带来的数值误差

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