はじめまして．インターン生の長谷川直哉と申します．こちらでは，私がNTT研究所での「画像強調モデルへのTest-timeAdaptationの適用検証」というテーマのインターンで取り組んだ研究と，そこから生まれた提案手法「ToSIE」について紹介させていただきます．
自己紹介
私は東京大学理学部情報科学科に所属している学生です．インターン参加時点では学部４年で，まだ研究室には配属されておらず (幣学部では研究室配属は４年の秋学期の初め)，研究をするのは本インターンが初めてでした．
以前から深層学習に興味があり，自分で深層学習モデルを構築してみたり，訓練させてみたりしていたのですが，なかなか実用に足るモデルができず，苦戦していたのを覚えています．この解決に役立つ知識である機械学習モデルのロバスト性について学部の授業で学ぶことはあったものの，なかなかそれを実用に生かす機会がなくこの時期まで過ごしてきました．その時に今インターンの募集を見つけ，実践を通してこれについて学ぶことのできる良い機会だと思い，応募させていただいた次第です．
テーマの背景について
深層学習モデルのロバスト性
深層学習モデルを含めた多くの機械学習モデルは，入力データの分布シフトに脆弱であるという弱点を持っています．
分布シフトとは学習時に使用したデータと，テスト・運用時に使用したデータの分布が一致しない現象を指します．例えば，自動運転における画像のセグメンテーションを考えてみましょう．訓練時は以下の図１のように晴れた昼の画像で訓練を行い，検証データで十分に高い精度の出力が得られたとします．
しかし，実用時はそう単純なことはありません．当然夜になったり，雨が降ったり，あるいはレンズが汚れて画像がぼやけてしまうかもしれません（図２）．このとき分布シフトが起こっており，このような汚染を受けた画像(破損したデータ)を入力された場合，モデルの出力は大きく劣化することがしばしばあります．
訓練時点でこのようなシフトが起こることがわかればよいですが，実際には予測不可能な分布シフトが起こることも十分に考えられます．このため，深層学習モデルを実用に足るものにするには，このような分布シフトに対しても頑健（ロバスト）なモデルを作る必要があります．
では，モデルのロバスト性を増す手法として既存手法にどのようなものがあるのか，今回は以下２つに絞って紹介します．
既存手法：画像強調 (Image Enhancement)
画像強調は画像を入力とするモデルのロバスト性を上げる手法の一つです．アイデアは至極単純で，破損した画像を認識モデルが認識しやすい画像に変換するというものです．この変換を行うモデルを作成し，推論時にあらかじめ画像をこのモデルに通すことで精度の悪化を防ぎます（図３）．
この手法の一つとして，URIEを紹介します．こちらは図４のようなU-Net型の構造を持つ深層学習モデルです．訓練時には破損を受けた画像を入力し，URIEの出力を認識モデルに入力し，タスク損失を最小化するようにパラメータを学習します．こちらも深層学習モデルなので，当然訓練していないタイプの破損を受けた画像が入力されたときは，一般に精度が落ちてしまいます．
既存手法：Test-time Adaptation (TTA)
Test-time Adaptation (以下TTAと呼ぶ)は画像強調とは異なり，モデル自体を入力データに適応させるように推論時に訓練も同時に行う手法を指します（図５）．また，この枠組みでは，事前訓練時に使ったソースデータセット，過去に入力されたテストデータセットは使えず，推論を行う際にそのミニバッチを一回だけ使って適応させ，そのデータを破棄するという制約があります．この制約は実世界ではストレージ容量の問題やプライバシーの問題（監視カメラのデータなど）でデータの蓄積が不可能という状況に対応します．
こちらの先行研究をいくつか紹介します．
バッチノーマライゼーション層の統計量をテスト時に適応させるだけのシンプルな手法です．今回はBNと呼びます．
分類モデルに使える手法です．今回はTentと呼びます．モデル出力を疑似ラベルとし，BN層をエントロピー最小化するように最適化する手法です．この手法の弱点として，大きな分布シフトにより疑似ラベルの精度が低下し，モデルが崩壊してしまうという点があります．
Tentを改良して連続的に変化する分布シフトにも対応させたモデルです．こちらはCoTTAと呼びます．
今回のテーマについて
以上２つの手法は競合することなく同時に適用可能です．つまり，図６のように画像強調モデルと認識モデルを１つの深層学習モデルととらえ，それに対してTTAをかけるという手法です．
この組み合わせにより，各手法について以下のようなメリットがあり，結果的にモデル全体ではより精度の高い推論ができると予想されます．
- 画像強調 … 未訓練のタイプの破損に対しても適切な修復をするように適応できる．
- TTA … Tent, CoTTAのような自己教師あり学習を行うモデルでは，疑似ラベルの精度が上がるため，より精度が向上する．
このことから本インターンでは，事前訓練した画像強調モデルと分類器をTest-time Adaptationでテストデータに適用することで，各々を単独で用いた場合より高い精度が得られるかを検証していきます．
予備実験
予備実験１
以上のアイデアを、画像強調手法単独での精度、TTA手法単独での精度と比較することで検証します。
実験では画像分類モデルの精度で比較を行います．今回使用した画像強調モデルはURIE，TTA手法はBN, Tent, CoTTAです．また，２つの手法を組み合わせる際，TTAを適用させる範囲が３通りあり，それぞれについても比較を行います．
- 分類モデルのみにTTA (Base adapt)
- URIEのみにTTA (URIE adapt)
- 両モデルにTTA (All adapt)
直感的には，モデルの一部にTTAを掛けるとその部分は現在の分布に適応するものの，ほかの部分が適応しないことから乖離が発生するため，全部にTTAを掛ける場合のほうがより高い精度を出すように思えます．
また，使用したデータセットはCIFAR100-Cの15corruptionsのSeverity 5です．このデータセットは15の異なるタイプの破損を受けた画像を含んでおり，各タイプの破損の画像がすべて入力されると別のタイプの破損の画像が入力されることで変化する分布シフトを再現しています．Severityは破損の度合いを示しており，5は最大の破損度を示しています．
Optimizerは学習率1e-5のAdam，認識モデルはAugMixで訓練されたResNeXt，URIEはCIFAR-100-Cの15corruptionsで学習済み¹です．
結果としてエラー率 (低いほど精度が高い) を以下の表１に示します．行方向に各手法，列方向に破損タイプ15個+全破損タイプの平均を示しています．なお，見やすさを考慮してCoTTAでの結果を省略していますが，Tentと同様な傾向が見られました．
全破損タイプで明らかに優れている手法はありませんが，平均だけ見てみますとTent, URIEだけを使う場合が優れているように見えます．組み合わせた手法(下3列)についてみてみますと，TTAの掛け方によらずほぼ同じような結果が出ているように見えます．このことから，学習率が小さすぎたのではないかと推測されます．よって，次は高い学習率で行った場合について実験してみます．
予備実験２
というわけで，学習率を1e-2に変え，同様な条件で実験を行ってみます（表２）．こちらも同様にCoTTAでの結果は省いております．
今回の場合，URIEのみにTTAを掛ける場合が全体的に見ても最もよくなりました．また，これとURIEだけの手法を比較してみて精度が向上していることから，URIEを現在の破損に適した修復をできるように適応させることに成功していることがわかります．
ただ，分類器の部分にTTA (Tent)をかける手法が大きなエラー率を出しています．これがTentに見られる誤差の蓄積で，間違った推論結果による自己教師あり学習が誤りを加速させた結果によるものです．学習率が高いことによりこれがより促進されているようです．
以上のことから，
- 分類器の部分は低い学習率で適応させたほうがよい
- URIEの部分は高い学習率で適応させたほうがよい．
ということがわかりました．今度はこのような条件設定で実験を行ってみます．
予備実験３
今度は分類器の部分の学習率を1e-4, URIEの部分の学習率を1e-2にして同様な実験を行います（表３）．
URIEとTTAを組み合わせた手法があまり低いエラー率を出していないことわかります．これについて破損タイプごとに分析してみましょう（表４）．
例えば，Gaussian noiseについては，URIE, TentはいずれもSourceよりも低いエラー率を達成しており，この２つを組み合わせた手法はより低いエラー率を出しています．対してImpulse noiseについては，URIEはSourceよりも高いエラー率を出しています．つまり，URIEを組み合わせないほうが良い結果を出すので，URIEとTentを組み合わせた手法がTentだけの手法よりも高いエラー率を出してしまう結果になるように推測されます．
このことから，URIEを使うかどうかを選択できるモデルを構築するのが良いのではないかと考えました．これが提案手法につながります．
本実験
提案手法：ToSIE
以上の実験から提案手法”Test-time adaptation of Selective Image Enhancement (ToSIE)”を考案しました．
モデルは図７のような全体図で，推論/適応時は以下のように動作します．
- URIEを通した画像とそうでない画像両方について認識モデルの出力２つを得る．
- ２つの出力のうち確信度の高い方を選択し，出力とする．
- 各TTAの手法に基づき，モデルを訓練する．
これにより，URIEを使うのが適していると考えられる場合のみ，URIEを通し，精度を高めることができます．ただ，URIEを通らない出力が採用された場合，このモデルをナイーブに訓練すると，URIEのパラメータが学習されません．今回の提案手法ではこの場合のみ，通常の損失に加え，URIEを通った出力を生徒ラベル，通らなかった出力を教師ラベルとするクロスエントロピー誤差を通常の損失に加えています．例えばTentの場合は以下のようになります²．
実験設定
提案手法を既存の手法，ナイーブにアンサンブルで２つの手法を組み合わせたモデルと比べることで有効性を確かめます．比較する手法は以下の通りです．
- URIE, 各TTA手法(Tent, CoTTA)のみ
- URIEを使うモデルと各TTA手法を使うモデル２つを用意して，それらをバギングしたモデル
- 提案手法ToSIE (Tent, CoTTA)
バギングは，ToSIEと同様に確信度の高いほうの出力を採用します．学習率はすべての手法で予備実験３と同じで，TTAはURIE，分類器両方にかけます．
結果
Tentに提案手法を適応したものは同様に誤差の蓄積が起きているものの，CoTTAに提案手法を適応したものは平均で30.90%と低いエラー率を達成しています (単独の手法よりも約5.5%, 単純なバギングよりも約2.5%の改善)．破損タイプ別にみても全体的に低いエラー率を達成しています．
まとめ
今回の研究は画像強調とTTAを同時に適用することでより良い精度を達成するモデルを作れるか検証するというものでした．これは以下の理由で，ナイーブに組み合わせるだけでは失敗に終わります．
- 画像強調, 分類器でTTAの適切な学習率が異なる
- 破損タイプによっては画像強調を適用しないほうが良い精度を達成する場合もある．
私はこの２つの問題点を解決するように構成した手法ToSIEを提案し，実際に高い精度を達成することを示しました．
補足：その他の検証
より一般的な評価を行うため，別のモデル，データセットでも検証を行いました．使用したモデルはResNet-50で，使用したデータセットはCUB-200です．CUB-200は画像サイズが224*224と大きく，また様々な物体の画像を含むCIFAR-100と異なり，鳥類画像のデータに限定されています．
その他の条件設定は予備実験３と同じです．結果としてエラー率のグラフを図８にまとめました．TTA手法別に色分けし，記号でTTAの掛け方を分けています．
URIEは全体的にSourceよりも改善しているものの，TTA手法を使った場合平均的にはSourceよりも悪化してしまうという結果が出ました．
非常に奇妙な結果ですが，観察してくると見えてくることがあります．例えばBNの結果（黄色のグラフ）に注目してみると，これがほかのTTA手法のエラー率のおよその下限になっていることがわかります．これはBNを含めたTTA手法の実装が原因であると考えられます．
そもそもBNはバッチノーマライゼーション層を訓練時と同じ挙動にすることで実装できます．今回の私の実験ではTent, CoTTAはCoTTAの論文の公式実装をもとにしており，この実装でもバッチノーマライゼーション層は訓練時と同じ挙動をするようになっています．つまり，この実装上ではTTA手法はBNを含むことになるため，BN自体がSourceよりも悪化してしまう今回の状況ではそれよりも改善させることが難しいのではないかと考えられます．
以上のことから，Tent, CoTTAについてバッチノーマライゼーション層を推論時と同じ挙動をするように変更した場合，精度がどう変化するかを実験してみた結果が図９になります．
以上のように，Tentでの誤差の蓄積を除けば，TTAを用いた場合のエラー率が確かにSourceよりも改善していることがわかります．以上のようにBNのエラー率が他TTA手法のそれの下限になってしまう場合はバッチノーマライゼーション層の挙動を推論時と同じように変更すればよいことがわかりました．
この変更は破損を受けていない画像での精度を上げるのにも役立つと考えられます．BNの実装では，ソースデータセットで訓練したバッチノーマライゼーション層の統計量を破棄してしまうため，Sourceよりも破損を受けていない画像での精度は下がってしまいます．バッチノーマライゼーション層を通常の推論時と同様にすることでこの現象は避けられると考えられます．
以上のように，BNで精度が悪化する場合はバッチノーマライゼーション層の挙動を変更することでTTA手法を用いて精度を向上させられることがわかりました．しかしながら，どのような条件でこのような現象が起きるかが不明確でした．CUB-200が鳥類画像に限定された狭い領域の画像のデータセットであるのに対してCIFAR-100が幅広い物体の画像データセットであること，バッチノーマライゼーション層の統計量の分布をみてみたところCUB-200で訓練したResNet-50にはいくつか外れ値が確認されたのに対しCIFAR-100で訓練したResNeXtにはそれが見られなかったことなど気になる差異はありましたが，検証不足なのでこれは将来の研究に託したいと考えております．
終わりに
研究の経験もなく，画像強調，TTAという手法はほぼ未知であった私ですが，何とか提案手法を残すことまででき，非常に達成感を感じています．途中，環境構築で躓いた私をサポートしてくださったり，進捗報告での議論，論文紹介等で知見を深めてくださったりと，特にメンターの榎本さん，担当の江田さんには深くお世話になりました．３週間という短いながらも非常に濃密な期間で，自分の能力を高めるよい機会であったと思います． そのほかインターンのテーマや環境の準備，研究および発表のアドバイスをしてくださったお二人をはじめとするNTT研究所の方々に深く感謝申し上げます．
今回の記事は以上となります．最後まで読んでくださりありがとうございました．
脚注
[1] 本来，学習とテスト時で同じ破損を使うべきではないが，今回は検証のために同じ破損を使いました．
[2]CoTTAの場合，損失はエントロピーではなく教師モデルのラベルを教師ラベルとするクロスエントロピーで学習するので，この損失でもそのラベルを教師ラベルとして用います．
