[論文レビュー] Crafting a Toolchain for Image Restoration by Deep Reinforcement Learning

Ke Yu, Chao Dong+2 more|arXiv (Cornell University)|Apr 10, 2018

Advanced Image Processing Techniques参考文献 41被引用数 31

ひとこと要約

本稿では、複雑で混合された歪みをもつ画像を復元するため、軽量でタスク特化型の畳み込みニューラルネットワークのツールボックスから動的に選択する強化学習フレームワーク、RL-Restoreを提案する。復元を段階的決定プロセスとして定式化し、段階的PSNRに基づく報酬とエージェントおよびツールの共同学習を採用することで、大規模CNNと同等の性能を達成しつつ、計算量を82.2%も削減し、未知のアーティファクトに対しても耐性が向上した。

ABSTRACT

We investigate a novel approach for image restoration by reinforcement learning. Unlike existing studies that mostly train a single large network for a specialized task, we prepare a toolbox consisting of small-scale convolutional networks of different complexities and specialized in different tasks. Our method, RL-Restore, then learns a policy to select appropriate tools from the toolbox to progressively restore the quality of a corrupted image. We formulate a step-wise reward function proportional to how well the image is restored at each step to learn the action policy. We also devise a joint learning scheme to train the agent and tools for better performance in handling uncertainty. In comparison to conventional human-designed networks, RL-Restore is capable of restoring images corrupted with complex and unknown distortions in a more parameter-efficient manner using the dynamically formed toolchain.

研究の動機と目的

• 単一の大規模CNNが、高いパラメータ数と計算コストを伴う複雑で混合された画像歪みを効果的に処理できないという限界に対処すること。
• 軽量でタスク特化型のネットワークの順次組み合わせを動的に実行することで、適応的かつパラメータ効率の良い復元を可能にすること。
• エージェントとツールの共同学習により、中間状態における不確実性やアーティファクトに対する耐性を高め、未知または未確認の歪みに対しても強靭な復元を実現すること。
• 可視的なツールチェーン選択により、段階的復元プロセスを明らかにすることで解釈可能性を向上させること。

提案手法

• フレームワークは、それぞれのタスク（ぼかし補正、ノイズ除去、JPEGアーティファクト低減）に特化した小さな専用CNNのツールボックスを用いる。
• エージェントは、現在の画像状態と過去の行動に基づき、段階的報酬関数（PSNR向上に比例）に従ってツールを逐次選択する方策を学習する。
• エージェントとツールを同時に学習する共同学習スキームにより、中間復元状態における不確実性やアーティファクトをよりよく扱えるようにする。
• エージェントは、画像品質が満足できると判断された段階でプロセスを自動的に停止でき、過剰適合を防ぎ、計算量の節約が可能である。
• 報酬関数はPSNR、SSIM、最終PSNR、最終MSEを用いて評価され、段階的PSNRがPSNRおよびSSIM両方の指標で最も優れた性能を示した。
• ツールの学習データには、合成ノイズと圧縮を含めることで、中間状態におけるアーティファクトに対する耐性が向上した。

実験結果

リサーチクエスチョン

• RQ1 強化学習エージェントは、複雑で混合された歪みをもつ画像を復元するため、軽量ネットワークの動的ツールチェーンを効果的に構成できるか？
• RQ2 エージェントとツールの共同学習は、未知または中間状態のアーティファクトに対して、どのように耐性を向上させるか？
• RQ3 異なる報酬関数の選択が、復元性能および収束特性に与える影響は何か？
• RQ4 自動停止機構の有無が、復元品質および計算効率に与える影響は何か？
• RQ5 本フレームワークは、単一の大規模CNNよりも、未確認の歪みに対して一般化性能が優れているか？

主な発見

• RL-Restoreは、82.2%のFLOPs削減を達成しながら、単一の大規模CNNと同等のPSNR性能を示し、顕著な計算効率の向上を実証した。
• 共同学習戦略により、ベースラインツールより0.2 dBのPSNR向上が得られ、さらに学習データにノイズ増強を組み込むと、追加で0.2 dBの向上が得られた。
• 段階的PSNR報酬関数は、段階的SSIM、最終PSNR、最終MSEと比較して、PSNRおよびSSIM両方の指標で優れた性能を示した。
• 自動停止機能を削除すると、平均で0.15 dBのPSNR低下が生じ、特に軽微な歪みに対してはより大きな低下が見られたことから、過剰復元のリスクが示された。
• 長さ12のツールチェーンと3つのツールが最適な性能を達成した。長くても性能の向上は顕著ではなく、学習の複雑さが増すためである。
• 本フレームワークは、解釈可能な段階的復元プロセスを明らかにし、ブラックボックスCNNに比べて透明性を提供した。