1. はじめに

こんにちは、黒糖ぱんです。

最近、グラフィックス界隈でよく耳にする「インバースレンダリング（画像から3Dシーンを逆算する技術）」や、 シェーダーアートでお馴染みの「レイマーチング」の仕組みに興味があり、自分でも一から実装して入門してみたいと考えていました。

インバースレンダリングのエンジンを自作するとなると、描画結果と正解画像の誤差を計算し、パラメータを最適化するための「自動微分（Autodiff）」の仕組みがほぼ必須になります。 普通、こういう機械学習的な最適化処理を書くのであれば、PyTorchなどを筆頭とするPythonエコシステムに乗っかるのが王道中の王道です。

……が、個人的に「どうしてもPythonを書きたくない」という強いこだわり（あるいはPythonアレルギー）がありました。

そこで、「Python以外で、モダンで使い勝手の良い自動微分フレームワークはないか？」と探していたところ、Rust製のディープラーニングフレームワーク「Burn」に行き着きました。 BurnはWGPUバックエンドを備えており、Rustから直接GPUリソースを叩きつつ、テンソル演算と自動微分の恩恵を受けられるという、今回の用途にぴったりの代物でした。

「RustでBurnを使ってレイマーチングを書けば、そのまま最適化計算に突っ込んでインバースレンダリングができるのでは？」

そんな思いつきから、実験的に「Burnを用いた微分可能レイマーチングエンジン」を作ってみました。本記事では、その実装の裏側や工夫したポイント、そしてそもそもインバースレンダリングとは何をやっているのかという基礎について紹介していこうと思います。

2. 作ったもの（デモと概要）

今回作成したプロジェクトの名前は、そのまま「Burn Raymarching」です。

百聞は一見に如かずということで、まずは「何ができるのか」を見てみましょう。

画像のように、最初はただのデタラメな球だったものが、学習が進むにつれてターゲットである三色団子の形へと変化していきます。

ざっくり何をしているのか？

一言でいうと、「複数の視点から撮った画像（2D）」を正解データとして与えると、それを再現するような「3D空間の球の集まり（位置・色・大きさ）」と「光の当たり方」を、オプティマイザ（Adam）が勝手に探し出してくれるシステムです。

全体のパイプラインとしては、以下の3つのステップで動くようになっています。

1. データセットの準備 (bin/generate) 学習の「正解」となる多視点の画像群（ターゲット画像）と、それぞれのカメラの位置情報（cameras.json）を用意します。今回は実験のため、自前のスクリプトで理想的なターゲット画像を生成しています。
2. 学習（Training） (bin/train) ここが今回のキモです。初期状態ではたった数個の適当な球を配置した状態からスタートします。Burnの自動微分（Autodiff）を使って、「現在の描画結果」と「正解画像」の誤差を計算し、球の位置や色、環境光の強さを少しずつ修正していきます。途中で不要な球を消したり（Pruning）、表現が足りない部分の球を分裂させたり（Splitting）しながら、指定した世代（Stage）分だけ学習を回します。
3. リアルタイムビューアー (bin/viewer) 学習が終わると、最適化された球のパラメータが scene.json に書き出されます。これを読み込み、RustのWGPUとWGSL（フラグメントシェーダ）を使って、学習結果をリアルタイムにグリグリ動かして見ることができるビューアーも作りました。

結果として、ポリゴンメッシュやボクセルすら使わず、純粋に球の集合だけで表現することができました。

「画像から3Dができるのは分かったけど、中の計算はどうなってるの？」と思うかもしれません。 次からは、この魔法のような処理を支えている「レイマーチング」の基礎と、それを「微分可能」にするための工夫について解説していきます。

3. 基礎知識：レイマーチングとインバースレンダリング

「画像から3Dができる」という結果を見たところで、その裏側で何が起きているのか、前提知識を少しだけ解説します。

レイマーチングとSDF

3D空間を描画する手法として最も一般的なのは、ポリゴン（三角形）を並べるラスタライズ手法ですが、今回はレイマーチング（Raymarching）を採用しています。

レイマーチングでは、SDF（Signed Distance Field：符号付き距離場）という数式を使います。これは、「空間の任意の座標 vec3(x, y, z) を入れると、そこから一番近い物体までの距離を返してくれる関数」です。 カメラから画面の各ピクセルに向かってレイ（光の筋）を飛ばし、「SDFが返す距離の分だけレイを進める」という処理を繰り返すことで、物体との交点を見つけて色を塗ります。

今回は、このSDFを使って「大量の球体」を空間に定義しています。

インバースレンダリングの仕組み

インバースレンダリングでは、この描画プロセスを微分可能（Differentiable）にします。

1. フォワードパス（描画）: 現在の球のパラメータ（位置・半径・色）と光源のパラメータを使って、特定のカメラ位置からレイマーチングで2D画像を生成する。
2. Loss（誤差）の計算: 生成した画像と、ターゲットとなる正解画像を見比べて、どれくらい違うか（Loss）を計算する。
3. バックワードパス（逆伝播）: Lossを減らすためには、各球の位置や色、光源を「どっちの方向に」「どれくらい」動かせばいいのか（＝勾配）を、自動微分を使って逆算する。
4. オプティマイザによる更新: Adamなどの最適化アルゴリズムを使って、パラメータを少しだけ更新する。

これを何千回も繰り返すことで、最初はデタラメだった球の集まりが、徐々に正解の3D形状へと収束していくわけです。

4. 実装のコア①：SDFと自動微分をどう繋ぐか

理屈は単純ですが、実際にこれを実装しようとすると「レイマーチングの不連続性」という大きな壁にぶつかります。ここが今回のプロジェクトで一番重要なポイントです。

ハードな描画は勾配を殺す

通常のレイマーチングは、空間のどこからでも距離が分かるSDFを使うため、一見すると微分の相性が良さそうに思えます。しかし、「レイが最初に衝突した物体の色をそのまま塗る（Zバッファ的な描画）」という処理をしてしまうと途端に学習が破綻します。

例えば、手前の球が奥の球を完全に隠している場合、画面には手前の球の色しか出力されません。すると誤差の責任（勾配）はすべて手前の球に押し付けられ、隠れている奥の球には「どう動けばいいか」という勾配が一切伝わらなくなります（勾配消失）。 また、物体の輪郭部分でも、ピクセルが「物体か背景か」で急激に色が切り替わるため、微分不可能な断絶が起きてしまいます。

解決策：Softmax関数による「半透明な」加重平均ブレンディング

そこで今回の renderer_diff.rs の実装では、レイが物体に衝突した際の「色の取り方」を工夫し、あえて境界を曖昧（Soft）にするアプローチをとりました。

レイの終点（交点）に対して、「一番近い球の色だけをポンッと取る」のではなく、「交点付近にある全ての球からの距離を計算し、近い球の色ほど強く反映されるように混ぜ合わせる（加重平均）」という処理を行っています。

// renderer_diff.rs より一部抜粋

// Step A: 交点と、全ての球の中心との距離を計算
let dists = dists_sq.clamp_min(1e-6).sqrt() - radius.transpose(); 

// Step B: 距離を元に、どの球の色をどれくらい使うか（重み）をSoftmaxで計算
let weights = softmax(dists.mul_scalar(-10.0), 1); 

// Step C: 各球の色に重みを掛けて足し合わせる（ブレンド）
// ★テンソルの次元合わせパズル（unsqueeze_dim）が発生する
let colors_expanded = colors.unsqueeze_dim::<3>(0);
let weights_expanded = weights.unsqueeze_dim::<3>(2);

let weighted_colors = colors_expanded * weights_expanded;
let mixed_color = weighted_colors.sum_dim(1).squeeze_dim(1);

この softmax を使ったブレンディングが魔法の役割を果たします。 球が少しだけ移動して交点との距離が変わると、それに連動して weights が滑らかに変化し、最終的な出力色も滑らかに変化します。この「滑らかな変化の繋がり」ができることで、初めて誤差の勾配が元のパラメータ（球の位置）まで正しく逆伝播されるようになります。

実は、深度を合成する際にも単なるZバッファ（一番手前のものを描画）ではなく、このような確率的・透過的なブレンディングを行うのが、最近の微分可能レンダラー（3DGSなど）の重要なテクニックだったりするらしいです。

5. 実装のコア②：球が増殖・消滅する適応的密度制御

インバースレンダリングをやる上で、最初から「空間のどこに、いくつの球を配置すれば正解か」なんて分かるわけがありません。 そこで、最近の3DGSの手法に倣って、学習の途中で動的にプリミティブ（球）の数を増減させる適応的密度制御（Adaptive Density Control）の仕組みを実装しました。

今回は全体の学習ステップをいくつかの「世代（Stage）」に分け、世代の変わり目でパラメータの状態を評価して以下の2つの処理を行います。

• Pruning（間引き） 最適化の過程で、物理的な半径が一定の閾値を超えて肥大化してしまった球や、黒くなりすぎて描画結果にほとんど影響を与えていない無駄な球をネットワークから削除します。
• Splitting（分裂） 世代間での「座標の移動距離」と「半径」を評価関数として見ます。学習を回しても位置がフラフラと大きく動いている球は、「その局所領域ではまだ表現の解像度が足りず、誤差の収束地点を見つけられていない」とみなせます。こういう球は、より小さな2つの球に分裂させてディテールを埋めにいきます。

最初はたった7個の球からスタートしても、この仕組みを回すことで、エッジなどの細かい表現が必要な場所には密集し、背景などの不要な場所からは球が消えていくようになります。

Burn特有の「テンソルサイズ変動」という壁

理屈は単純なのですが、これをBurn（というか静的な計算グラフを構築するディープラーニングフレームワーク全般）で実装しようとすると、「テンソルサイズの変動でOptimizerが爆発する」という厄介な問題にぶち当たります。

AdamなどのOptimizerは、過去の勾配の移動平均（モメンタム）などの内部状態を、最適化対象のテンソルと同じサイズ（今回なら球の数 N×3 など）で保持しています。そのため、Stageの途中でPruningやSplittingが走って球の数 N が変わると、テンソルの形状が不一致を起こして演算が即死します。

これを回避するため、train.rs では以下のような泥臭いライフサイクル管理を行っています。

1. 世代交代のタイミングで新しい球の数とパラメータを計算し、モデルのテンソルを作り直す。
2. 古いAdamの内部状態（モメンタム等）は諦めて破棄し、新しいテンソルサイズに合わせてAdamConfigから完全に初期化・再構築する。
3. Adamをリセットした反動で学習が吹っ飛ばないよう、Stageが進むにつれて学習率（Learning Rate）を段階的に減衰（アニーリング）させ、後半はより局所的な微調整に注力させる。

「モデルごとOptimizerを作り直して勾配を繋ぎ直す」というかなり力技なアプローチですが、これによりAutodiffの計算グラフを破綻させることなく、必要な箇所へ動的に計算リソースを割り当てることに成功しています。

6. おまけ：複雑な形を学習させようとしてAdamが爆発した話

「三色団子ができたなら、もっと複雑な形（DNAヘリックスやドーナツ）もいけるのでは？」と欲を出して試してみたのですが、現実は甘くありませんでした。

球の重なり（オクルージョン）が増えたり、配置が複雑になったりすると、Adamオプティマイザが簡単に局所解（Local Minima）に陥ってしまい、謎のゴミの塊が生成されたり、球が描画範囲外にすっ飛んでいったりします。 インバースレンダリングにおいて「初期値の配置」や「学習率の繊細な調整」がどれほど重要か（そして、それを力技で解決している最近の3DGSなどの論文がどれほどヤバいか）を身をもって知りました。むずすぎではこれ。

7. まとめ：テンソル演算の限界と真の「今後の展望」

最終的に、「理想的なデータ」であれば、Adamオプティマイザの力でターゲットの3D形状とライティングを復元することには成功しました。

今回、RustのDLフレームワーク「Burn」を使ってテンソル演算のゴリ押しで微分可能レイマーチングを実装しましたが、やってみて痛感したのは「行列処理のパズルでレイマーチングを書くのは苦行」ということです。 全レイと全球の距離計算をバッチ処理するために unsqueeze_dim や matmul で必死に次元を合わせていると、直感的なシェーディングコードが全く書けなくなります。

もし今後、こういった「微分可能レンダリング」をガチで研究・実装するなら、最近グラフィックス界隈で注目を集めている Slang のような「微分可能シェーダー言語」を使うのが圧倒的に正解だと思います。あちらなら、普段のHLSLのような文法で書きつつ、コンパイラレベルで自動微分の恩恵を受けられるため、今回のような行列パズルに苦しめられることはありません。

機械学習とグラフィックスの境界線は年々曖昧になってきており、ブラックボックスの中身を自作して学ぶのは非常にエキサイティングです。 レンダリングや変態的な最適化アルゴリズム、あるいはLoss設計とパラメータ調整という終わりのない沼に興味がある人は、ぜひこの世界に足を踏み入れてみてください！