損失関数の性能を上げる手法SAMの要約とdual norm problem

# 導入

『[SoTAを総なめ！衝撃のオプティマイザー「SAM」爆誕&解説！](https://qiita.com/omiita/items/f24e4f06ae89115d248e)』を読んでいると「dual norm problemがわからない」という記述があったのですが、その説明がコメントで書くには長くなりそうだったので記事にしました。
文脈であるSAMの紹介をした後に、dual norm problemとは何かを説明し、dual norm problemを解きます。

# SAMの超要約

損失関数が与えられたとき、近傍での損失の最大値を新たな損失関数として定義する手法がSAMです。

要約を書いていたところ長くなってしまったので短くしましたが、せっかくなので元のバージョンも残しておきます。

# SAMの要約

SAMとはある損失関数を元に新しい損失関数を作る方法、及びその新しい損失関数の勾配の近似計算法のことです。
$L^2$ 正則化はある損失関数に対して $L^2$ 損失を加えた損失関数を考えるという手法ですが、SAMも概ねこれと同じジャンルに属します。
この要約では新しい損失関数の作り方の部分を紹介します。

損失関数 $L(w)$ が与えられているとします。ここで$w$はモデルのパラメータです。[^1]
[^1]: 教師あり学習において損失は教師データとモデルのパラメータを用いて与えられるものですが、ここでは教師データは関数$L$に織り込まれているとしています。
SAMでは新たな損失関数は次の式で与えられます。

```math
L^\text{SAM}(w) = \max_{\|\epsilon\| \le \rho} L(w + \epsilon) + \lambda \|w\|_2^2
```

$\rho, \lambda > 0$ はハイパーパラメータで $\\|\epsilon\\|$ は $\epsilon$ のノルムです。SAMでは特に $L^p$ ノルムを利用するようです。
右辺第二項は単なる $L^2$ 正則化であり、第一項がSAMの核となる部分です。
第一項は $w$ での新たな損失として $w$ の周辺での損失の最大値を考えるということです。
これは周辺で一様に損失が少ないところを探す、つまり安定して $L(w)$ が小さくなるパラメータを探すということを意味し、直感的にもっともらしい手法です。

# dual norm problem

$\newcommand{\argmax}{\mathop{\rm arg~max}\limits}
\newcommand{\sign}{\mathop{\rm sign}}
\newcommand{\trans}{\mathrm{T}}
\newcommand{\ip}[2]{\left\langle#1,\,#2\right\rangle}$
$L^\text{SAM}(w)$ の定義にある $\max$ を実現する $\epsilon$ を $\hat{\epsilon}$ と書くことにすると、$\hat{\epsilon}$ は $L(w + \epsilon)$ を $\epsilon$ について一次近似することにより次のように近似できます。(詳しい話は[元記事](https://qiita.com/omiita/items/f24e4f06ae89115d248e)を参照してください。)

```math
\hat{\epsilon}(w) \approx \argmax_{\|\epsilon\| \le \rho}\,\ip{\epsilon}{\nabla L(w)}
```

ここで $\ip{\epsilon}{\nabla L(w)}$ は $\epsilon$ と $\nabla L(w)$ の内積です。
このままでは記号にノイズが多いので次の一般的な形で考えることにします。

```math
\argmax_{\|y\| \le \rho}\,\ip{y}{x}
```

内積の線形性から次の等式が成り立ちます。

```math
\argmax_{\|y\| \le \rho}\,\ip{y}{x} = \rho\ \argmax_{\|y\| \le 1}\,\ip{y}{x}
```

この右辺の $\rho$ より後ろの部分を求める問題、つまり $x$ を用いて $\ip{y}{x}$ を最大化する $y$ を計算する問題を指して論文はdual norm problemと呼んでいると思われます。
なぜこれをdual norm problemと呼ぶかという話は後回しにして、まずこの問題の解が元の記事で書かれているような形になることを導出します。以下ではノルムは $L^p$ ノルムとします。

まずは $\max_{\\|y\\|\_p \le 1}\, \ip{y}{x}$ を計算してみましょう。
鍵となるのはヘルダーの不等式[^2]という不等式です。
$p,q \ge 1$ が $1/p + 1/q = 1$ を満たすとき、次の不等式が成り立ち、この不等式をヘルダーの不等式と呼びます。

[^2]: 証明は「[数学対話: ヘルダーの不等式とコーシー・シュワルツの不等式](http://aozoragakuen.sakura.ne.jp/taiwa/taiwaNch04/totu/node4.html)」を参照してください。

```math
\ip{y}{x} \le \|y\|_p \|x\|_q
```

議論が複雑になるのを避けるために以下 $p,q > 1$ とします。
この不等式から $\\|y\\|\_p \le 1$ のとき次の不等式が成り立つことがわかります。

```math
\ip{y}{x} \le \|y\|_p \|x\|_q \le \max_{\|y\|_p \le 1}\, \|y\|_p \|x\|_q = \|x\|_q
```

これで $\ip{y}{x} \le \\|x\\|\_q$ がわかったので、あとは実際に $\\|y\\|\_p \le 1$ かつ $\ip{y}{x} = \\|x\\|\_q$ を満たすような $y$ を構成できれば、それが $\argmax\nolimits_{\\|y\\|\_p \le 1} \ip{y}{x}$ であるということになります。
そこで内積を取ったときに $|x\_i|$ の $q$ 乗が現れるように $y\_i = \sign(x\_i) |x\_i|^{q-1}$ と定義して内積を計算してみます。
ここで $\sign(x\_i)$ は $x\_i$ の符号であり、 $\sign(x\_i)\,x\_i = \|x\_i\|$ を満たすことに注意しましょう。

```math
\ip{y}{x} = \sum_i \sign(x_i) |x_i|^{q-1} x_i = \sum_i |x_i|^q = \|x\|_q^q
```

右辺は $\\|x\\|_q$ になってほしかったので両辺を $\\|x\\|\_q^{q-1}$ で割り、次の等式を得ます。

```math
\ip{\frac{y}{\|x\|_q^{q-1}}}{x} = \|x\|_q
```

$1/p + 1/q = 1$ の両辺に $q$ を掛けて適当に移項することで $q - 1 = q/p$ が得られるので、$y / \\|x\\|\_q^{q-1} = y / \\|x\\|\_q^{q/p}$ です。
$y / \\|x\\|\_q^{q/p}$ の $L^p$ ノルムが $1$ 以下であれば、これを新たに $y$ として取り直すことで目当ての $y$ が構成できたことになります。
$y / \\|x\\|\_q^{q/p}$ の $L^p$ ノルムを計算するために、まず $y$ の $L^p$ ノルムを計算してみましょう。
$q-1 = q/p$ より $y_i = \sign(x_i) |x_i|^{q/p}$ であることに注意します。

```math
\begin{align}
\|y\|_p &= \left( \sum_i \left| \sign(x_i) |x_i|^\frac{q}{p} \right|^p \right)^\frac{1}{p} \\
&= \left( \sum_i |x_i|^q \right)^\frac{1}{p} = \left[ \left( \sum_i |x_i|^q \right)^\frac{1}{q} \right]^\frac{q}{p} \\
&= \|x\|_q^\frac{q}{p}
\end{align}
```

よって $y / ||x||\_q^{q/p}$ の $L^p$ ノルムは次のように計算できます。

```math
\left\| \frac{y}{\|x\|_q^{q/p}} \right\|_p = \frac{\|y\|_p}{\|x\|_q^{q/p}} =  1
```

ゆえに $\argmax$ が $y / \\|x\\|_q^{q/p}$ であるとわかりました。
成分を計算すると元記事と同じ形になります。

```math
\frac{y_i}{\|x\|_q^{q/p}} = \sign(x_i) \frac{|x_i|^{q-1}}{(\|x\|_q^q)^{1/p}}
```

# dual norm problemの数学的背景

$\newcommand{\R}{\mathbb{R}}$
ここから先は本題に関係ありません。
なぜ上述の問題がdual norm problemと呼ばれるかがわかる内容になっています。

$\R^n$ から $\R$ への一次関数全体の集合を $\R^n$ の双対空間 (dual space) と呼び、$\R^{n*}$ と書くことにします。
$\R^n$を縦ベクトルの集合だと思ったとき、一次関数は横ベクトル$a$を用いて $f(x) = ax$ と書くことができるので、双対空間は横ベクトルの抽象的な定義と解釈することができます。
$\R^n$ 上では $L^p$ ノルムなど種々のノルムを考えることができますが、$\R^n$ 上のノルムに応じて $\R^{n*}$ にもノルムを定義したいです。
例えば一次元の場合、$f(x) = ax$ という一次関数には $|a|$ という値をノルムとして対応させることが自然に思われますが、これは変換の倍率になっています。
そこで一般に最大の倍率という意味合いで次のように双対空間上のノルムを定義し、双対ノルム (dual norm) と呼びます。

```math
\|f\| = \max_{\|x\| \le 1} f(x)
```

これは $\\|x\\| \le 1$ という球体上での $f(x)$ の最大値になっており、元のノルムが変われば球体の形状が変わるので双対ノルムの値も変化します。

ところで、ベクトル $x \in \R^n$ に対して $I_x(y) = \ip{y}{x}$ という一次関数を対応させることで$\R^n$から$\R^{n*}$への対応ができます。
これは行列の転置に対応し、この対応により $\R^n$ と $\R^{n*}$ を同一視できます。つまり $\R^n = \R^{n*}$ ということです。
ノルムも込めて考えると $(\R^n, \\|\cdot\\|)^* = (\R^n, \\|\cdot\\|_*)$ となります。ただし右辺のノルムは $\\| x \\|\_\star = \\|I\_x\\|$ で定義します。
そこで双対ノルムは実際に計算するのが難しい形だったので、双対ノルムの値と双対ノルムの定義の $\max$ を実現する$y$を具体的に決定するという問題 (dual norm problem) が生まれます。

$L^p$ ノルムの場合は上で行った計算がその解答になっていて、$1/p + 1/q = 1$ とすると

```math
\|x\|_{p*} = \|I_x\|_* = \max_{\|y\|_p \le 1} \ip{y}{x} = \|x\|_q
```

であり、$\max$ を与える $y$ は

```math
y_i = \sign(x_i) \frac{|x_i|^{q-1}}{(\|x\|_q^q)^{1/p}}
```

となることがわかります。ゆえに $(\R^n, \\|\cdot\\|\_p)^* = (\R^n, \\|\cdot\\|\_q)$ です。
実はこの話は無限次元を含む一般的な主張の特別な場合であり、一般的な場合は関数解析という分野で取り扱われます。関数解析を学ぶと双対空間がどう役立つかも知ることができます。