みんなのためのAI 第5回 — 47万パラメータの U-Net で白黒写真に色をつける、そしてなぜ色が褪せたのか

• はじめに — 出力が画像のとき
• U-Net — skip connection が運ぶもの
• 学習
• 結果 — 形は合っていて色は褪せた
• なぜ色が褪せたのか — L1 損失の性質
• 実際のカラー化モデルはどうしているか
• まとめ
• 🧠 理解度チェッククイズ
• 参考資料

はじめに — 出力が画像のとき

これまでは出力がテキストでした。第1回は文、第3回は単語、第4回はキャプションです。今回は 出力そのものが画像 です。

課題はカラー化です。白黒写真を入れると色をつけて返します。

入力: 32×32 白黒 (チャンネル1個)
出力: 32×32 カラー (チャンネル3個)

データを別途集める必要はありません。CIFAR-10 のカラー画像を白黒に変えれば 入力と正解のペアが自動的にできあがります 。

rgb = torch.stack([...])                                    # 元のカラー = 正解
gray = (rgb * torch.tensor([0.299, 0.587, 0.114])
        [None, :, None, None]).sum(1, keepdim=True)         # 白黒変換 = 入力

重み 0.299, 0.587, 0.114 は、人の目が緑に最も敏感で青に鈍いことを反映した標準的な輝度式です。こうして作ったデータを 自己教師あり学習 (self-supervised) と呼びます。人がラベルを付けていないのに正解が存在するからです。

U-Net — skip connection が運ぶもの

カラー化には厄介な要求があります。 「何であるか」を知って色を決める必要がある一方、「どこにあるか」も正確でなければなりません。 カエルだと分かって緑を塗るわけですが、カエルの輪郭の外に緑がはみ出してはいけません。

一般的なエンコーダ・デコーダはこの両方を同時に満たすのが難しいです。エンコーダが画像を圧縮する過程で意味は得られても位置情報を失うからです。U-Net はこれを skip connection で解決します。

class TinyUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.e1 = block(1, 32); self.e2 = block(32, 64); self.e3 = block(64, 128)
        self.d2 = block(128 + 64, 64); self.d1 = block(64 + 32, 32)
        self.out = nn.Conv2d(32, 3, 1)

    def forward(self, x):
        s1 = self.e1(x)                       # 32x32  — 最も細かい
        s2 = self.e2(F.max_pool2d(s1, 2))     # 16x16
        h  = self.e3(F.max_pool2d(s2, 2))     # 8x8    — 最も抽象的
        h = self.d2(torch.cat([F.interpolate(h, scale_factor=2), s2], 1))
        h = self.d1(torch.cat([F.interpolate(h, scale_factor=2), s1], 1))
        return torch.sigmoid(self.out(h))

核心は torch.cat([アップサンプルされたh, s2], 1) です。デコーダが解像度を戻すとき、 同じ解像度のエンコーダ出力をそのまま持ってきてつなげます。

• 下から上がってきた h — 低解像度だが「これはカエルだ」といった意味を含む
• 横から来た s2 — 圧縮を経ていないので輪郭と質感が無傷

デコーダはこの2つを一緒に見ます。意味は深い経路から、位置は skip 経路から来ます。チャンネル数が 128 + 64 になっているのはこのためです。

sigmoid で終わるのも意図的です。出力が0〜1の範囲の RGB 値でなければならないので、活性化で範囲を強制します。

学習

opt = torch.optim.AdamW(model.parameters(), lr=2e-3)

while not run.over_budget():          # 12分の予算
    ix = torch.randint(0, len(rgb), (128,))
    x, y = gray[ix].to(dev), rgb[ix].to(dev)
    loss = F.l1_loss(model(x), y)
    opt.zero_grad(); loss.backward(); opt.step()

損失は L1 (平均絶対誤差) です。予測した RGB 値と正解の RGB 値の差をそのまま測る回帰損失です。この選択が結果の性格を決めるのですが、後ほど改めて扱います。

[   14.1s] pairs=30000
[   14.2s] params {"params": 472323, "res": 32}
[  722.0s] SUMMARY {"steps": 62136, "final_l1": 0.0231}

62,136ステップ、最終 L1 は0.023。ピクセル値が0〜1の範囲なので 平均誤差が2.3% という意味です。数字だけ見れば見事です。

結果 — 形は合っていて色は褪せた

各行は3枚ずつ2組で、順序は 入力の白黒 / モデルの予測 / 正解のカラー です。

正直に読んでみましょう。

うまくいったこと — 形が正確に保存されています。猫の毛、船の甲板構造、カエルの脚、自動車の窓枠がぼやけていません。skip connection が役目を果たした証拠です。色の 方向 も概ね合っています。カエルの周りは緑がかり、船の周りは青みがかり、猫は茶系に向かいました。

うまくいかなかったこと — 色が全体的に 褪せています 。正解では鮮やかな赤の自動車が、予測では暗い灰褐色です。正解の赤い船体と青い海が、予測では濁った灰青色です。セピア調に近い結果が多くなっています。

なぜ色が褪せたのか — L1 損失の性質

これはモデルが小さいからではありません。 損失関数を L1 に選んだ結果 です。

白黒の自動車写真を1枚考えてみましょう。この車は赤かもしれず、青かもしれず、白かもしれません。白黒の情報だけでは決められません。つまり 正解が1つではなく複数ある問題 です。

このとき L1 損失はモデルにこう言います。「正解との絶対誤差を最小化せよ」。複数の正解があり得る状況でこの指示に従う最も安全な方法は何でしょうか。

あらゆる可能性の中間値を打つこと です。

赤を打って正解が青なら大きな罰点を受けます。しかし灰色を打てば、正解が何であっても罰点は中程度に抑えられます。学習が進むほどモデルはこの安全な戦略に収束し、その結果が 彩度の低い色 です。

L1 が0.023という良い数字と褪せた色が同時に成立する理由がこれです。 L1 はまさに言われたことをよくやったのであり、私たちが望んだものがそれではなかっただけです。

実際のカラー化モデルはどうしているか

この問題はよく知られており、解法もいくつかの系統に分かれています。

色を分類問題に変える — Zhang らの2016年の研究は、色空間を313個の区間に分け「このピクセルはどの区間か」を当てる分類として解きました。分類では複数の候補に確率を分けられるので平均でつぶれません。サンプリングのときに鮮やかな色を選べばよいのです。

敵対的損失を加える — GAN の判別器は「この画像は本物らしいか」を見ます。褪せた色は本物の写真らしく見えないので罰点を受けます。L1 が形を、判別器が鮮やかさを担当する組み合わせです。

知覚損失を使う — ピクセル値を直接比較する代わりに、事前学習された神経網の特徴空間で比較します。人が感じる類似性により近くなります。

3つに共通点があります。 平均に逃げることが得にならないように損失を設計する という点です。

まとめ

47万パラメータで形を無傷に保ちながら色をつけました。skip connection のおかげで細かい構造が生き残り、自己教師あり方式なのでラベリング費用は0でした。

そしてこの回の本当の教訓は色が褪せた側にあります。 損失関数は目標を定義するものであって、目標を達成する道具ではありません。 L1 を選んだ瞬間に「平均に近い答えが安全だ」というルールも一緒に決まり、モデルはそのルールに忠実に従いました。結果が気に入らないなら、モデルではなく 何を最小化せよと言ったのか をまず振り返るべきです。

次回はこの問題を正面から扱う方式 — 拡散モデル を見ます。複数の正解があり得る状況で平均に逃げず、1つを選び取る構造です。

🧠 理解度チェッククイズ

参考資料

• CIFAR-10 データセット
• U-Net — 生物医学画像分割のための畳み込みニューラルネットワーク
• Colorful Image Colorization — 色を分類として解くアプローチ
• Image-to-Image Translation with Conditional Adversarial Networks (pix2pix)