LDTを動かしてみる

# 概要
2024年4月に公開された論文で提唱されている機械学習モデルの話
- [Latent diffusion transformer for point cloud generation](https://link.springer.com/article/10.1007/s00371-024-03396-1)

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3363422/11434270-5b25-455a-b5fc-495be1d9c726.png)

詳しくは[LDT概要](#LDT概要)に記載。


# 動作環境
||Version|
|---|---|
|Ubuntu|20.04|
|python|3.10.12|
|Graphics Board|NVIDIA GeForce RTX 4090|


# データセット
ShapeNetCore.v2.PC15k


# インストール手順
- [こちら](https://github.com/Negai-98/LDT)からリポジトリをクローン
- 必要モジュールをインストール
```bash
pip install -r requirements.txt
```
- 必要ファイルをダウンロード
[Drive](https://drive.google.com/drive/folders/1FRRKDBFNQTW_HdDglNro4ufJVPsD8zcz)内の`extern`をダウンロードし、`LDT`直下に配置
- セットアップ
```bash
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
```
- データセットの用意
```bash
./dataset_download.sh
```
でダウンロード可能
`LDT`直下に`data`ディレクトリを生成し、`ShapeNetCore.v2.PC15k`を格納する
これで一通り準備は完了


# 実行
基本的に以下3ステップで学習完了となる
その後学習モデルを使って生成する
1. Compressorの学習
1. 学習モデルをセッティング
1. Diffusion Generatorの学習

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


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

しかしあるメソッドへ引き渡される変数が足りないなどコード上の問題が生じていたため、以下に自身が修正した内容を記載しておく
該当箇所は`trainer/Compressor_Trainer.py`29~41行目
```python
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`の該当箇所に記載する
```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`部分
```yaml
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
![Screenshot from 2025-04-25 14-12-09.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3363422/78c0f0af-ec83-4af0-9bdd-94c0022fd8c7.png)

### 4000-6000
![Screenshot from 2025-04-25 13-39-01.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3363422/177f1956-673f-4a4f-a8a0-79328e69b264.png)
![Screenshot from 2025-04-25 13-43-07.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3363422/d5edaf83-1dbc-4755-8c91-c4b01c468dcc.png)

## chair

### 4000-6000
![Screenshot from 2025-04-25 14-29-15.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3363422/b8f3fab6-c5c4-460f-977d-fa0b5d3d90d0.png)

---

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

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

---
---

# LDT概要
LDTは[こちら](https://mail.google.com/mail/u/0/?tab=rm&ogbl#chat/space/AAAATjAr7j8)の論文に詳しく書かれている。本記事の画像もこちらから引用している。

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3363422/11434270-5b25-455a-b5fc-495be1d9c726.png)

LDTは大きく分けて以下の２つで構成されている
- [Latent Compressor](#Latent-Compressor)
- [Diffusion Generatorの学習](#Diffusion-Generator)

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


## Latent Compressor

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

---

### 構成要素

#### 1. Downsampling

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

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

- **入力**：点群の座標データ
```math
X \in \mathbb{R}^{N \times3}
```
- **出力**：クラスタ特徴ベクトル
```math
X_0 \in \mathbb{R}^{M \times D}
```
---
```math
C = \text{FPS}(X)\\
```
```math
X_0 = \text{MaxPooling}(\text{MLP}(\text{KNN}(X, C)))
```

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

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

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

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

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

#### 損失関数

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

```math
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)$ の学習に利用**される
- 各ステップでは、以下のようにガウスノイズを加えていく

```math
  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$ を復元**する


```math
  p_\theta(Z_0 | Z_T) = p(Z_T) \prod_{t=1}^{T} p_\theta(Z_{t-1}|Z_t)
```
```math
  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 によって処理し、**各層ごとの平均・分散を動的に調整**

```math
  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)$ をキーやバリューに使用

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

#### 損失関数

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

```math
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$：時間ステップごとの重み

---


## データの流れ
以下の表にまとめる  

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