Beyond scalar losses: calibrating segmentation models via gradient vector field surgery

あなたは、ロボットに人間の体内の隠れた都市の地図を描く方法を教えていると想像してください。このロボットは、「ディープラーニングモデル」と呼ばれる特別な「魔法の目」を使い、X線やMRIのような医療スキャンを見て、腫瘍や血管といった小さく困難なものを見つけ出します。目標は、ロボットがこれらの箇所を見つけるだけでなく、自分が何を見ているのかについて、どの程度確信を持っているかを知ることです。もしロボットが「これは腫瘍であると99%の確信を持っています」と言ったなら、それが正解でなければなりません。もしそれが単なる影であったなら、ロボブルは「あまり自信がありません」と言うべきなのです。

長い間、これらの小さな箇所を見つけるためにロボットを教える最良の方法は、「Dice loss（ダイス損失）」と呼ばれるルールを使うことでした。このルールは、ゲームのスコアのようなものです。もしロボットが描いた形が実際の形と完璧に重なれば、高いスコアが得られます。これは、画像の中で対象が非常に小さい場合でも、形を見つける上では非常にうまく機能します。しかし、落とし穴があります。このルールを使うと、ロボットは形を見つけることには長けますが、自分の自信の度合いを知ることに関してはひどい状態になります。つまり、自信過剰な「自信満々な愚か者」になってしまうのです。間違っているときでも「100%確信しています！」と叫んでしまうのです。現実の世界において、医師が手術の判断にこのロボットを使用する場合、自信満々に間違えるロボットは危険です。それは、晴れている日でも必ず「雨が降ります」と言い張る天気予報士のようなもので、その助言を信頼することができません。

「Beyond scalar losses: calibrating segmentation models via gradient vector field surgery」という題名のこの論文は、この過剰な自信の問題に取り組んでいます。ドイツと米国の研究チームである著者らは、ロボットが学習する方法（具体的には、間違いに基づいてどのように「脳」を調整するか）こそが、その傲慢さを生んでいることを発見しました。彼らは、「gradient vector field surgery（勾配ベクトル場の手術）」と呼ぶ巧妙な解決策を提案しました。単にゲームのルールを変えるのではなく、ロボットの学習プロセスに対して、非常に精密で小さな操作を行います。彼らは、ロボットが間違いに対してどのように反応するかを微調整することで、形を見つける能力を失うことなく、自身の自信についてより謙虚かつ正確にさせることができると発見しました。彼らは、眼や乳房の2D画像から、脳や腎臓の3Dスキャンに至るまで、さまざまな医療画像を用いてテストを行い、彼らの手法がロボットの予測をはるかに信頼できるものにすることを示しました。

問題点：自信過剰なロボット

なぜロボットがこれほど図に乗ってしまうのかを理解するには、それがどのように学習しているかを見る必要があります。想像してみてください。ロボットは、画像内の小さなピクセルが腫瘍（前景）の一部であるか、それとも単なる正常組織（背景）であるかを推測しようとしています。ロボットは、「腫瘍である確率は70%」といった推測をします。すると、教師（損失関数）が答えをチェックします。

これらのタスクにおける標準的な教師は、「Dice loss」を使用します。この教師は主に「形」を重視します。もしロボットが腫瘍をうまく覆うような塊を描けば、教師は満足します。しかし、ここが奇妙な点です。Dice教師は、ロボットが「どのように」その形に到達したかには関心がありません。ロボットがデタラメに推測していたのか、あるいは確信を持っていたのかは気にしないのです。

この論文は、Dice教師には奇妙な癖があることを示しています。それは、もしロボットが非常に自信を持っていても（例えば99%の確信があっても）間違っていた場合、学習させる必要はほとんどないと判断してしまうことです。ロボットの「学習信号（勾配）」がほぼゼロになってしまうのです。これは、問題を間違えたにもかかわらず、本人が「自分は正しい」と強く信じ込んでいるために、教師が「この生徒は自信がありすぎて話を聞かない」と判断して、修正をやめてしまう学生のようなものです。逆に、もしロボットが確信を持てず（50/50の状態）、迷っているときは、教師は変化するように激しく叫びます。これにより、奇妙なループが生じます。ロボットは、答えを極端な方向（0%または100%）へと押し出すことを学習してしまうのです。なぜなら、そこでは教師が修正を放棄するからです。その結果、ロボットは形を描くことには非常に精緻になりますが、たとえ間違えていても病的なまでに自信満々になってしまうのです。

解決策：Gradient Vector Field Surgery

著者らは、自信の問題を解決するためには、単にロボットが得るスコアを変えるだけでは不十分であることを悟りました。学習信号の「方向」そのものを変えなければならないのです。彼らはこれを「gradient vector field surgery（勾配ベクトル場の手術）」と呼んでいます。

ロボットの学習プロセスを、谷の底（完璧な解）を探しているハイカーだと想像してください。「勾配（グラディエント）」は、ハイカーにどちらの方向に進むべきかを教える地面の傾斜です。従来のDice教師を用いた場合、傾斜は端の方で平坦で混乱しており、ハイカーをマップの極端な角（答えが0または1になる場所）へと押しやってしまいます。

著者らの「手術」は、傾斜に特別な因子を加えることを含みます。彼らは、学習信号が誤差に応じて線形にスケールするように調整します。このように考えてみてください。もしロボットが少し間違っていれば、小さな後押しを与えます。もし大きく間違っていれば、大きな突きをします。しかし決定的なのは、彼らが「ブレーキ」を追加し、ロボットが速すぎる自信を持つのを防ぐようにしたことです。

彼らは、ハイカーのためのカスタムマップとして機能する新しい数学的公式を設計しました。このマップは以下のことを保証します：

1. ロボットはエラーから学ぶ： もし間違っていれば、たとえどれほど自信を持っていたとしても、信号は強力になります。
2. ロボットは謙虚さを保つ： 真に確信を持たない限り、ロボットが答えを0%や100%へと押し切ってしまうのを防ぎます。

彼らがこれを「手術」と呼ぶのは、学習プロセスの数学的な流れの中に切り込み、新しい挙動を縫い合わせるからです。興味深いことに、彼らはこの新しい「マップ」が、従来のDiceやCross-Entropy（交差エントロピー）のような、単一の単純なスコア（スカラー損失）では記述できないことを発見しました。それはより複雑で多次元的な流れです。これは数学者にとっては恐ろしいこと（「非保存的」な場がロボットをぐるぐる回らせることを懸念するため）ですが、著者らは、特定の医療画像においては、ロボットは回転することなく、より良く、より正直な解へと真っ直ぐに進むことを証明しました。

結果：信頼できる予測

チームは、網膜血管の2D画像、乳房の腫瘤、および脳転移や腎臓腫瘍の3Dスキャンを含む、多様な医療データセットに対してこの手術をテストしました。彼らは、彼らの手法を標準的なDice loss、Cross-Entropy loss、およびその他の一般的な組み合わせと比較しました。

結果は驚くべきものでした。彼らが標準的なDice lossにこの「手術」を適用したところ：

• 自信が劇的に向上した： 「期待較正誤差（Expected Calibration Error）」（自信がどれほど間違っているかの指標）が大幅に減少しました。例えば、INbreastデータセットでは、誤差は0.0058から0.0026に減少しました。FIVESデータセットでは、0.0162から0.0044に減少しました。
• 精度は高いまま維持された： ロボットは形を見つける能力を失いませんでした。実際、いくつかの困難なデータセットでは、精度（Dice Similarity Coefficientで測定）は向上するか、あるいは同等を維持しました。例えば、INbreastデータセットでは、Diceスコアは64.93%から74.34%へと跳ね上がりました。
• どこでも機能する： 2Dまたは3D画像で使用した場合でも、あるいはTverskyやCombo Lossのような他の損失関数と組み合わせた場合でも、この手術は一貫してモデルの較正（キャリブレーション）を向上させました。

著者らはまた、彼らの公式にある指数パラメータ $n$ という「つまみ」についてもテストしました。彼らは、$n$ を上げていく（$n$ を大きくしていく）につれて、パフォーマンスが向上し、$n=20$ 付近でプラトー（停滞）に達することを発見しました。これは、彼らの手法が堅牢であることを意味します。完璧にチューニングするために数学の魔術師である必要はありません。

なぜこれが重要なのか

この論文は、この「勾配手術（gradient surgery）」が、過剰な自信の問題を解決するためのシンプルかつ強力な方法であると結論づけています。ロボットが間違いから学ぶ方法を変えることで、彼らは腫瘍を見つけるだけでなく、自分がどれほど確信しているかについても正直になれるモデルを作り上げました。

これは臨床への導入において非常に大きな意味を持ちます。もし医師が手術中の切除部位を決定するためにAIを使用しているなら、AIが推測しているのか、それとも確信を持っているのかを知る必要があります。「腫瘍であると99%の確信を持っていますが、間違っているかもしれません」と言うロボットは、実際に間違っているときに「100%の確信があります」と言うロボットよりも、はるかに有用です。著者らは、このアプローチが、あらゆる医療用セグメンテーションネットワークをより信頼できるものにするための一般的なツールとなり、ヘルスケアにおけるより安全で信頼できるAIへの道を切り開く可能性があると示唆しています。彼らは単に新しい損失関数を見つけたのではありません。学習プロセスそのものをより誠実にする方法を見つけたのです。