LiteKD-Net: Lightweight Knowledge-Distilled Network for Mobile Image Denoising

技術要約：モバイル画像デノイジングのための LiteKD-Net

問題提起
モバイル画像のデノイジング（ノイズ除去）は、デバイス上での展開に向けて、高い復元品質の達成と計算コストの低減という二重の課題に直面しています。モバイルデバイスはプロフェッショナル向けカメラよりも小さなセンサーと画素を使用するため、本質的にS/N比（信号対雑音比）が低く、単純な加法性ガウスモデルでは捉えきれない複雑で信号依存的なノイズパターン（例：フォトンのショットノイズ、リードノイズ、行ノイズ）が発生します。さらに、このタスクのための深層学習モデルの学習は、厳密に整合された低品質（LQ）画像とグラウンドトゥルース（GT）画像のペアの不足によって阻害されます。これは、異なるノイズ特性を持つ同一のシーンを撮影することが、動きやカメラ内処理の変動により困難であるためです。既存の高容量ネットワーク（数千万のパラメータを持つものなど）は、モバイルハードウェアにおける高レイテンシ、過剰なエネルギー消費、および大きなメモリフットプリントに悩まされています。また、クラウドベースの処理は、レイテンシとプライバシーの問題を引き起こします。

手法
著者は、データの不足、モデルの効率性、および性能の回復に対処するために、3つのコアコンポーネントを統合した統一フレームワークである LiteKD-Net を提案しています。

1. ノイズシミュレーション・パイプライン (NSP): 整合された実世界のペアの欠如を克服するため、著者は（もともとDSLRカメラ向けであった）物理ベースのノイズモデルを、クリーンなsRGB画像から合成訓練データを生成するために適応させました。このパイプラインは、sRGB画像を線形化し、フォトンショットノイズ、Tukey–Lambdaリードノイズ、および行ノイズをシミュレートし、極めて重要な点として、ピクセルクロストークを導入します。このクロストークは、隣接する画素間に局所的な空間相関を生み出す特定の畳み込みカーネルを介してモデル化され、コンパクトなモバイルセンサーに見られる電気的/光学的な漏れを模倣します。最後に、量子化ノイズが加えられ、画像はクリッピングされた後、sRGBに変換されます。
2. Lite-RRDB アーキテクチャ: 著者は、アップサンプリングモジュールを除去することで、Real-ESRGANアーキテクチャをアイデンティティ解像度のデノイジング用に適応させています。軽量な**Student（生徒）ネットワークを作成するために、標準的な残差密集ブロック（RRDB）を Lite-RRDB に置き換えました。この設計では、標準的な畳み込みが深度分離畳み込み（Depthwise followed by Pointwise）**に置き換えられ、密な特徴量の連結が削除され、固定チャネル幅（64）を用いた逐次処理へと変更されています。これにより、残差内残差（residual-in-residual）の訓練構造を維持しつつ、パラメータ数と計算複雑性を大幅に削減しています。
3. Noise-KD (知識蒸留): 軽量なStudentの容量減少を補うために、特徴レベルの知識蒸留戦略が採用されています。まず、高容量の**Teacher（教師）**ネットワーク（適応されたReal-ESRGAN）が学習され、その後固定されます。Studentの学習中、Studentは平均二乗誤差損失（$L_{KD}$）を用いて、Teacherの中間特徴マップに一致するように誘導されます。Studentは、ピクセルレベルの再構成損失（$L_1$）、知覚損失（$L_{per}$）、および蒸留損失を用いて共同で最適化されます。重要なことに、Teacherは学習後に破棄されるため、推論時の追加コストは発生しません。

主な貢献

• ノイズシミュレーション・パイプライン (NSP): ピクセルクロストークを明示的にモデル化することで、モバイル特有のノイズシミュレーションのギャップに対処し、完全に整合されたGTおよびLQデータを生成する新しいパイプライン。
• Lite-RRDB モジュール: 標準的な畳み込みを深度分離畳み込みに置き換え、密な連結を削除することで、モデルサイズとFLOPsを劇的に削減する軽量なアーキテクチャブロック。
• Noise-KD 戦略: 固定された重いTeacherから軽量なStudentへ、推論パラメータや計算量を増やすことなく復元能力を転送する、特徴レベルの蒸留手法。
• 包括的な評価: 実世界のデータセットでモデルを評価し、復元品質と計算効率の間の強力なトレードオフを実証。

実験結果
実世界のデータセット（Mobile AI Denoising Dataset および MIDD）を用いた実験では、LiteKD-Netが計算メトリクスを大幅に削減しながら、競争力のある画像品質を維持していることが示されました。

• 効率性: ベースラインのReal-ESRGANと比較して、LiteKD-Netはモデルサイズを約10倍（16.70Mから1.67Mパラメータへ）削減し、MACsおよびFLOPsを4倍削減しました。
• 速度: 推論実行時間は、256×256の入力に対して約2倍（148msから83msへ）短縮され、その結果、より高いフレームレート（11.98 FPS vs. 6.72 FPS）を実現しました。
• 品質: 軽量なStudent単体では重いTeacherと比較して品質の低下が見られますが、Noise-KDの追加により性能が回復します。LiteKD-Netは、重いReal-ESRGANのベースラインに非常に近いPSNRおよびSSIMスコアを達成し、効率性の指標においてはSwinIRを上回っています。知覚的指標（LPIPS, DISTS, MUSIQ）に関しては、比較されたモデルの中で最高のMUSIQスコアを獲得するなど、競争力のある性能を発揮しています。

意義と主張
本論文は、LiteKD-Netがモバイル画像デノイジングにおいて復元品質と計算効率の優れたトレードオフを提供すると主張しています。物理学に基づいたノイズシミュレーション・パイプライン、軽量なアーキテクチャ、および知識蒸留を統合することで、大規模なネットワークによる過大なコストやクラウド処理によるレイテンシを回避しつつ、高品質なデノイジングモデルをモバイルデバイスに展開できることを著者は実証しています。結果は、提案手法が、モバイルハードウェアの厳格なリソース制約を遵守しながら、重いモデルの復元品質を効果的に回復できることを示しています。著者は、自らのシミュレーションが真のRAWデータではなく線形化されたsRGB画像上で動作しているため、パイプラインが完全なカメラ・イメージング・チェーン（例：デモザイキング、トーンマッピング）をモデル化しているわけではないと謙虚に述べていますが、sRGBベースのデノイジングのための整合された訓練データを生成する上で、依然として効果的です。