概要

2024年4月に公開された論文で提唱されている機械学習モデルの話

Latent diffusion transformer for point cloud generation

LDTはLatent Diffusion Transformerの略であり、その名の通り拡散モデルとTransformerを用いた生成aiモデルのことである。主な生成対象は3DPointCloud。

詳しくはLDT概要に記載。

動作環境

	Version
Ubuntu	20.04
python	3.10.12
Graphics Board	NVIDIA GeForce RTX 4090

データセット

ShapeNetCore.v2.PC15k

インストール手順

こちらからリポジトリをクローン
必要モジュールをインストール

pip install -r requirements.txt

必要ファイルをダウンロード
Drive内のexternをダウンロードし、LDT直下に配置
セットアップ

cd evaluation/pytorch_structural_losses/
python setup.py install
cd extern/pointnet2_ops_lib
python setup.py install
cd extern/emd
python setup.py install

データセットの用意

./dataset_download.sh

でダウンロード可能
LDT直下にdataディレクトリを生成し、ShapeNetCore.v2.PC15kを格納する
これで一通り準備は完了

実行

基本的に以下3ステップで学習完了となる
その後学習モデルを使って生成する

Compressorの学習
学習モデルをセッティング
Diffusion Generatorの学習

各学習機構のパラメータ数は以下の通り
Compressor:8100217
Diffusion Generator:457012344

1. Compressorの学習

3D点群データを拡散モデルで扱いやすい形式にするために情報を圧縮する装置
この段階では入力データを圧縮して展開するだけ
python3 train_Compressor.pyで学習できる。

しかしあるメソッドへ引き渡される変数が足りないなどコード上の問題が生じていたため、以下に自身が修正した内容を記載しておく
該当箇所はtrainer/Compressor_Trainer.py29~41行目

def update(self, data):
    cates = data["cate_idx"].to("cuda")
    data = data['tr_points'].to(self.device)
    self.model.train()
    self.warm_up(self.optimizer, self.itr)
    loss, kl_loss, rec_loss = self.compute_loss(data, cates)
    self.optimizer.zero_grad()
    loss.backward()
    if self.cfg.opt.grad_norm_clip_value is not None:
        clip_grad_norm_(self.model.parameters(), self.cfg.opt.grad_norm_clip_value)
    self.optimizer.step()
    self.itr += 1

    ####################これを追加#####################
    with torch.no_grad():                          
        model_output = self.model(data)
        max_feature = model_output['max'].max().item()
    #################################################
    
    return loss, kl_loss, rec_loss, max_feature     # max_feature追加

問題なく動くと./experiments/Compressor_Trainer/airplane/に学習モデルのファイルが保存されているはず

2. 学習モデルをセッティング

1で出力した学習モデルのファイルパスをexperiments/Latent_Diffusion_Trainer/airplane/config.yamlの該当箇所に記載する

compressor:
  pretrain_path:experiments/Compressor_Trainer/airplane/checkpt_5.pth #ここに追加
  outsize: 2048
  ...

3. Diffusion Generatorの学習

圧縮されたデータを拡散、逆拡散し、その過程を学習することで完全なノイズから意味のあるデータを生成するためのモデルを学習する装置
圧縮機と同様python3 train_Latent_Diffusion.py.pyで学習可能

ここでも変数不足のエラーが生じたため、修正内容を記載する
該当箇所はexperiments/Latent_Diffusion_Trainer/airplane/config.yamlのscore部分

score:
  num_steps: 1000
  z_dim: 120
  z_scale: 32
  hidden_size: 1024
  num_heads: 16
  num_blocks: 24
  num_categorys: 1
  c_dim: 0.
  t_dim: 1024
  dropout: 0.
  norm: layer_norm
  learn_sigma: False
  act: swish
  unet: False
  AdaLN: True
  condition: False
  graphconv: False #追記

これで実行が終わると./experiments/Latent_Diffusion_Trainer/airplane/に学習モデルが生成されているはずである。

また、その他にsmp_epXXX.npyというファイルがあり、これは訓練の段階で1つだけ生成された点群データである

結果

以下の組み合わせで学習させてみた

category	Compressorの学習epoch	Diffusion Generatorの学習epoch
airplane	5	6000
airplane	4000	6000
chair	4000	6000

airplane

5-6000

4000-6000

chair

4000-6000

かなりの高精度で生成できていることが分かる。
たまに歪なデータ(上記ID:airplane_40)が生成されることもあるが、ShapeNetデータセットの中にもおおよそ飛行機と呼べないようなデータもいくつか存在していたため、学習自体は問題ないと考えられる

今後

各モデルのパラメータなどを変化させ、データ生成の傾向を調査する。
epoch数を増やし過学習の影響を見る
→そのためには数値として評価指標を立てるべきだが、そのあたりも調べていく

LDT概要

LDTはこちらの論文に詳しく書かれている。本記事の画像もこちらから引用している。

条件付き生成と無条件生成の２種類が用意されているが、今回は無条件生成の方で話を進めていく

LDTは大きく分けて以下の２つで構成されている

Latent Compressor
Diffusion Generatorの学習

学習する際は上記2つをそれぞれ2回に分けて行う

Latent Compressor

目的

元の点群を、拡散モデルで扱いやすい潜在トークンへ変換する
複雑な点群形状 → 高表現力・柔軟なトークンへ圧縮

構成要素

1. Downsampling

点群を小さなクラスタに分割し、効率的に特徴を抽出
画像処理の「パッチ分割」に類似した手法を点群に適用

手順：

FPS（Furthest Point Sampling）：代表点（クラスタ中心）を抽出
KNN：各代表点の近傍点をクラスタ化
MLP + MaxPooling：クラスタごとの特徴を抽出

入力：点群の座標データ

X \in \mathbb{R}^{N \times3}

出力：クラスタ特徴ベクトル

X_0 \in \mathbb{R}^{M \times D}

C = \text{FPS}(X)\\

X_0 = \text{MaxPooling}(\text{MLP}(\text{KNN}(X, C)))

\begin{align}
X&：学習用入力データ\\
C&:クラスタ中心の座標データ\\
M&：クラスタ数\\
D&：1つのクラスタに含まれる点群の数
\end{align}

2. Tokenization Auto-Encoder

エンコーダー

Transformer ベースの構造
クラスタ中心を位置エンコーディングとして使用
自己注意（Self-Attention）により、局所的・大域的な依存関係を学習
点群の幾何情報をエンコーディングに含めることで、各トークンが異なる領域の特徴を効果的に表現
入力：クラスタ特徴ベクトル

X_0 \in \mathbb{R}^{M \times D}

出力：潜在トークン

Z \in \mathbb{R}^{M \times D} \

デコーダー

Cross-Attention を備えた Transformer
シードベクトル $h \in \mathbb{R}^{N \times D}$をクエリとして使用
クロスアテンションにより、トークン $Z$ の情報をシードに統合
最終的に MLP を通じて点群を再構成

再構成の特徴

シードからサンプルサイズ（クエリの数）を変えることで、様々な解像度の点群を再構成可能
潜在空間で明示的に幾何構造を保存しなくても、各トークンが局所領域を表すことで、ローカル情報を効果的に学習できる

損失関数

潜在トークン $Z$ を通じて点群 $X$ を再構成するための損失関数

L(\xi, \psi) = \mathbb{E}_{q_\xi(Z|X)}[-\log p_\psi(X|Z)] - \lambda_{\text{KL}} \cdot D_{\text{KL}}(q_\xi(Z|X) \parallel p(Z))

$\xi$：エンコーダのパラメータ
$\psi$：デコーダのパラメータ
$q_\xi(Z|X)$：エンコーダによる潜在トークンの分布
$p_\psi(X|Z)$：デコーダによる再構成点群の分布
$p(Z)$：標準正規分布
$\lambda_{\text{KL}}$：KLダイバージェンス項の重み
再構成誤差には Chamfer Distance (CD) や Earth Mover's Distance (EMD) を使用

Diffusion Generator

目的

Latent Compressorで変換された潜在トークンを用いて生成のための学習を行う
ノイズから潜在トークンを生成できるようにする

構成要素

1 Forward Diffusion（順拡散）

順拡散は マルコフ連鎖として設計されており、元の潜在トークン $Z_0$ に 段階的にノイズを加えるプロセスであり、ノイズ予測ネットワーク $\epsilon_\theta(Z_t, t)$ の学習に利用される
各ステップでは、以下のようにガウスノイズを加えていく

  Z_t = \sqrt{\alpha_t} Z_0 + \sqrt{1 - \alpha_t} \varepsilon_t, \quad \varepsilon_t \sim \mathcal{N}(0, I)

2 Reverse Diffusion（逆拡散）

逆拡散は、拡散生成における 生成プロセスそのもの
順拡散でノイズが加えられたデータ $Z_T$ を 段階的にノイズ除去して $Z_0$ を復元する

  p_\theta(Z_0 | Z_T) = p(Z_T) \prod_{t=1}^{T} p_\theta(Z_{t-1}|Z_t)

  p_\theta(Z_{t-1}|Z_t) = \mathcal{N}\left(Z_{t-1}; \mu_\theta(Z_t, t), \sigma_t^2 I\right)

このとき、平均 $\mu_\theta$ はノイズ予測ネットワークを用いて計算される
学習済みのノイズ予測ネットワーク $\epsilon_\theta$ を用いて、正規分布から徐々に潜在トークンを生成できる

3 Dynamic Transformer（動的トランスフォーマー）

拡散プロセスで用いる ノイズ除去ネットワークの中核
Transformer をベースとし、長距離依存性のモデル化が可能
さらに、時間ステップごとに変化するノイズ量に対応するために Dynamic Positional Encoding（動的位置エンコーディング） を導入
時刻 $t$ に基づく位置エンコーディング $P$ を MLP によって処理し、各層ごとの平均・分散を動的に調整

  Z^{(l+1)}_t = \text{AdaLN}\left(\text{MLP}\left(\text{softmax}\left(\frac{Q^{(l)}_t {K^{(l)}_t}^\top}{\sqrt{D^{(l)}}}\right)V^{(l)}_t\right), P_t\right)

条件付き生成（Conditional Generation）の場合は、外部条件 $Y$（例：画像や不完全な点群）から得た特徴 $\tau(Y)$ をキーやバリューに使用

  K_t = \tau(Y)W_K, \quad V_t = \tau(Y)W_V

損失関数

ノイズ付き潜在状態からのノイズ予測を通じて、点群を生成できるように訓練

L(\theta, \xi, \psi) = \mathbb{E}_{q_\xi(Z|X)}[-\log p_\psi(X|Z)] + \mathbb{E}_{q_\xi(Z|X)}[\log q_\xi(Z|X)] + \mathbb{E}_{q(Z_t|Z)}\left[w_t \cdot \left\| \varepsilon - \varepsilon_\theta(Z_t, t) \right\|^2\right]

$\theta$：拡散生成器（denoising network）のパラメータ
$\varepsilon_\theta(Z_t, t)$：ネットワークによるノイズ予測
$\varepsilon$：実際のノイズ
$w_t$：時間ステップごとの重み

データの流れ

以下の表にまとめる

フェーズ	動作	学習しているか？
エンコーダー	点群を潜在表現に圧縮	している
順拡散	ノイズを加える	していない
逆拡散	ノイズ除去してデータを推定	している
デコーダー	潜在表現から点群を復元	している