1. はじめに
先日公開した記事のレビュー中、プロファイリングに基づいた高速化の経緯とか詳しく知りたいというコメントをもらいました。
https://techblog.adacotech.co.jp/entry/2026/03/19/083000
そこで今回は、Rustで実装していたPCA処理を題材に、実際にボトルネックを見つけて改善していった過程を書いてみます。
なお、以下のコードは説明のために一般化・簡略化しています。
2. 対象コード
今回対象にしたのは、n << d なデータに対してPCAを行う処理です。
n: サンプル数d: 特徴量次元数
画像をそのまま特徴量ベクトルとして扱う場合、d はかなり大きくなります。
このとき、通常のPCAのように共分散行列 X^T X を作ると、d x d の行列を扱うことになり、計算量・メモリともに重くなります。
この問題を回避するために、固有値問題の解き方を変えるアプローチがあります。
PCAでは本来、(中心化済みデータに対する)共分散行列 X^T X の固有値問題を解きます:
X^T X v = λ v
ここで u = X v とおくと:
X X^T u = λ u
が成り立ちます。
重要なのは、0以外の固有値 λ は共通であるという点です(厳密には、min(n, d) 個の非ゼロ固有値が両者で共通になります)。
つまり、X^T X の固有値問題は、X X^T の固有値問題に対応しており、
後者を解いた結果から元の主成分ベクトルを復元できます。
X^T X: d x d(特徴量空間)X X^T: n x n(サンプル空間)
したがって、n < d の場合には、より小さい n x n の行列であるGram行列を用いることで、計算量とメモリを削減できます。
今回のコードでは、この Gram行列 G = X X^T を使ったPCA(dual PCA) を採用しています。
全体の流れは次のようになります。
use ndarray::{Array1,Array2}; struct PcaModel { components:Array2<f64>, } fn apply_pca(features:&Array2<f64>,target_dim:usize) ->Array2<f64> { // フル成分で学習 let full_dim=features.nrows(); let model=fit_gram_pca(features,full_dim); let transformed=transform(features,&model); // 最後に必要な分だけ切り出す transformed.slice(s![..,0..target_dim]).to_owned() } fn fit_gram_pca(features:&Array2<f64>,target_dim:usize) ->PcaModel { let gram=build_gram(features); let (eigen_values,eigen_vectors)=solve_eigen(&gram); let components= restore_components(features,&eigen_values,&eigen_vectors,target_dim); PcaModel {components } } fn build_gram(features:&Array2<f64>) ->Array2<f64> { features.dot(&features.t()) } fn transform(features:&Array2<f64>,model:&PcaModel) ->Array2<f64> { model.components.dot(&features.t()).t().to_owned() }
この実装は動作としては正しいのですが、実際には不要な計算を含んでおり、無駄が発生しています。
ここから、このコードに対して改善を入れていきます。
3. 不要な主成分計算を削減する
まず着目したのは、必要以上の主成分を計算している点です。
対象コードでは、以下のようにフル次元で学習したあとに切り出しています。
// 対象コード(再掲) fn apply_pca(features:&Array2<f64>,target_dim:usize) ->Array2<f64> { let full_dim=features.nrows(); let model=fit_gram_pca(features,full_dim); let transformed=transform(features,&model); transformed.slice(s![..,0..target_dim]).to_owned() }
しかし、最終的に必要なのは target_dim 次元だけです。
したがって、最初からその分だけ計算すれば十分です。
// 改善後 fn apply_pca(features:&Array2<f64>,target_dim:usize) ->Array2<f64> { let model=fit_gram_pca(features,target_dim); transform(features,&model) }
この変更により、
- 学習時の計算量
- モデルサイズ
- transform時の計算量
をまとめて削減できます。
この改善だけでも大きな効果がありました。
- 着手前:
188.659s - 初期改善後:
86.094s
ただし、この時点でもまだ十分に速くはありませんでした。
4. パフォーマンスをみる
4.1 まずは簡易計測から始める
どこが重いのかを知るために、まずは std::time::Instant を使った簡易計測から始めました。
use std::time::Instant; fn timed<T>(label:&str,f:impl FnOnce() ->T) ->T { let started= Instant::now(); let result = f(); eprintln!("{}: {:.3} ms",label,started.elapsed().as_secs_f64()*1000.0); result }
PCA周辺を次のように分解して計測しました。
let gram=timed("gram", ||build_gram(features)); let (eigen_values,eigen_vectors)=timed("eigen", ||solve_eigen(&gram)); let components=timed("restore", || { restore_components(features,&eigen_values,&eigen_vectors,target_dim) });
4.2 実際に測ると、重いのはGram行列計算だった
計測の結果、処理の内訳は次のようになりました。
| 処理 | 比率 |
|---|---|
| Gram計算 | 約85% |
| 主成分復元 | 約14% |
| 固有値分解 | 約1% |
PCAといえば固有値分解、というイメージから、最も計算コストが高いのは固有値分解だと考えていました。
が、しかし実際には、支配的だったのはGram行列の構築(X X^T)でした。
この時点で、改善対象をGram行列計算に絞ることにしました。
5. Gram行列計算を改善する
まず、対象コードのこの部分を改善していきます。
// 対象コード fn build_gram(features:&Array2<f64>) ->Array2<f64> { features.dot(&features.t()) }
5.1 対称性を利用して上三角だけ計算する
Gram行列は対称行列なので、上三角だけ計算すれば十分と考えました。
ndarrayの dot() はBLASフィーチャを有効化していない環境では適応的な補助関数に任せられるため、手で上三角のみ計算することで計算量を半減できるメリットが期待できます。
// 改善後(上三角のみ計算) fn build_gram(features:&Array2<f64>) ->Array2<f64> { let n=features.nrows(); let mut gram= Array2::zeros((n,n)); for i in 0..n { let row_i=features.row(i); for j in i..n { let value = row_i.dot(&features.row(j)); gram[(i,j)]=value; gram[(j,i)]=value; } } gram }
この変更は大きく効きました。
- 初期改善後:
86.094s - 上三角化後:
57.453s
約33%の短縮です。
5.2 ブロック化してキャッシュ効率を上げる
さらに、キャッシュ効率を改善するためにブロック化を行いました。
// 改善後(ブロック化) fn build_gram(features:&Array2<f64>,block_size:usize) ->Array2<f64> { let n = features.nrows(); let mut gram = Array2::zeros((n,n)); for row_start in (0..n).step_by(block_size) { let row_end= (row_start+block_size).min(n); let row_block=features.slice(s![row_start..row_end, ..]); for col_start in (row_start..n).step_by(block_size) { let col_end= (col_start+block_size).min(n); let col_block=features.slice(s![col_start..col_end, ..]); let block=row_block.dot(&col_block.t()); gram.slice_mut(s![row_start..row_end,col_start..col_end]) .assign(&block); if row_start!=col_start { gram.slice_mut(s![col_start..col_end,row_start..row_end]) .assign(&block.t()); } } } gram }
この変更でさらに改善しました。
- 上三角版:
57.453s - ブロック版:
54.796s
5.3 採用しなかった案
固有ベクトルの上位成分のみを求める部分分解も試しましたが、
- ブロックGram単体:
54.796s - 部分分解込み:
54.948s
と改善は見られませんでした。
理論的によさそうでも、実測で勝てなければ採用しない、という判断になります。
6. 生成AIによる最適化
これまでは、上記のような計測→ボトルネックを特定→改善案の適用→効果検証のサイクルを適用してきました。
最近では、CodexやClaude Codeを活用し
apply_pcaのプロファイル計測を行って下さい。 ボトルネックとなる箇所を特定し改善案を提案してください。
などと指示すると改善案を出してくれたりします。
改善施策適用後の効果検証もレポート作成まで任せる事が可能で、改善に要する時間は大幅に短縮できます。
7. まとめ
今回の最適化で一番重要だったのは、思い込みではなく実測でボトルネックを特定したことでした。(自分が予想したボトルネックとずれてました。)
改善の流れは次の通りです:
- 不要な主成分計算を削減
- 計測でボトルネックを特定
- Gram行列計算を最適化
最終的な処理時間は以下の通りです。
| 段階 | 時間 |
|---|---|
| 着手前 | 188.659s |
| 初期改善後 | 86.094s |
| 上三角化後 | 57.453s |
| ブロック化後 | 54.796s |
今回の範囲では、
- 数学的に不要な計算を減らすこと
- 行列計算をCPUに優しい形にすること
の2つが特に効果的だったという結論でした。