1️⃣

可変品質での圧縮を実現する画像トークナイザ「One-D-Piece」を公開しました

2025/01/22に公開

Turingの基盤AIチームの三輪です。

Turingで研究・開発した結果をまとめた「One-D-Piece: Image Tokenizer Meets Quality-Controllable Compression」という論文を公開したので、紹介します。

https://turingmotors.github.io/one-d-piece-tokenizer/

https://arxiv.org/abs/2501.10064

これは何?

「One-D-Piece」という新たな「画像トークナイザ」を提案しました。これによって256段階の可変品質で画像を圧縮し、Transformerベースのモデルと相性のいいトークン列での表現に変換することができます。

背景

テキストトークナイザと画像トークナイザ

Transformer型のモデルは、「トークン」の列を入力に取ります。そのため、Transformer型のモデルを用いるためには、データを何らかの方法でトークンにする「トークナイザ」が必要です。

トークナイザといって一般に想起されるのはテキストトークナイザ、つまりテキストをテキストトークンに分解するシステムです。こうしたテキストトークナイザはByte-Pair Encodingのようなアルゴリズムを用いてテキストをトークン列へと変換します。トークン列を「デトークナイズ」することで元のテキスト列を復元することができます。こうした仕組みはLLMをはじめとするテキストを扱う機械学習モデルで一般的に用いられています。

一方、テキスト以外のものもトークン列に変換できれば、Transformerアーキテクチャ上で容易に扱うことができます。特に、画像をトークン列に変換する仕組みを「画像トークナイザ」と言います。

画像トークナイザの振る舞いは、テキストトークナイザのそれと同様に考えることができます。何らかの方法で画像を分解し、トークン列を得ます。このトークン列をデトークナイズすることで元の画像を得ることができます。したがって、この中間表現のトークン列が画像を表現していると考えることができます。

画像トークナイザを考える上で難しいのは、その実際の仕組みです。「各ピクセルを1トークンとみなす」というような単純なアイデアはMNISTに含まれるような小さな画像を扱う上では悪くない選択肢ですが、少し大きい画像サイズ、例えば256x256の画像で65,536トークンものサイズになるのは後続タスクにおいて相当に不便です。さらに、1ピクセルはRGBの3バイトで表現されるため、このようなトークナイザの語彙は既に16,777,216個ということになります。通常の言語モデルにおける語彙はどんなに大きくても数十万であることを踏まえると、このような画像トークナイザは非常に扱いづらいものになるだろうと考えられます。

画像トークナイザの仕組み

そこで、多くの画像トークナイザは「劣化を許して情報を圧縮しつつ、十分な再構成後の品質を確保する」ことを目標とします。特に画像生成を目的とした場合には、AutoEncoder型の構造による圧縮と量子化を組み合わせるVQ-VAEの方式が主流です。

https://arxiv.org/abs/1711.00937

離散画像トークナイザの構造

離散化され、有限の語彙を持つ画像トークナイザは自己回帰による生成モデル(言語モデルなど)やMaskGITなどの離散拡散モデルに基づいた画像生成へのアプローチと相性が良い傾向にあります。Turingが開発する世界モデル「Terra」は自己回帰による生成モデルで、MAGVIT-v2という離散画像トークナイザが用いられています。

https://zenn.dev/turing_motors/articles/6c0ddc10aae542

Terraの大きな課題として「画像トークン長」が挙げられます。Terraに用いている画像トークナイザMAGVIT-v2は(288, 512)サイズの画像を(16, 16)のサイズのグリッドに分割し、それぞれの領域(画像パッチ)をトークナイズする構造です。このため、576(=18x32)トークンが画像1枚ごとに必要になります。従って、例えばTerraによってFPS=20で動画を3秒分生成する場合、34,560(=576x20x3)トークンを生成しなければなりません。このトークン長をどうにか圧縮できれば、生成を高速化することができます。

なお、トークナイザの「圧縮率」には様々な考え方があり、その1つは上のように単純にトークン長で見る方法です。一方、このトークナイザは262,144(=2^18)の語彙を持ちます。この場合、1トークンを表現するのに18bitを要することから、画像全体では1,296(=18x(32x18)/8)バイトであると主張することができます。このあとの議論では、このようにトークンをバイトに換算して比較することがあります。

TiTok

「トークン長の圧縮」という観点で、私たちは圧縮率が高いことで知られるTiTokというアーキテクチャに注目しました。

ByteDanceの研究者が提案する「TiTok」はNeurIPS 2024に採択された「An Image is Worth 32 Tokens for Reconstruction and Generation」という論文で導入されたものです。Vision Transformerに基づく新しいトークナイザアーキテクチャを用いることにより、良好な再構成性能と生成モデル性能を実現しつつ、32トークンという極めて小さいサイズに画像を圧縮することを実現しました。

https://arxiv.org/abs/2406.07550

従来のVQ-VAEのようなアーキテクチャでは、画像をパッチ単位に分割し、それぞれのパッチに対応するトークンを生成する手法が主流でした。しかし、この方法にはいくつかの課題がありました。隣接するパッチ間で高い相関を持つ情報をうまく統合できないことや、情報量が密集した領域で適切にトークナイズできないことが挙げられます。

TiTokはこれらの課題を克服するために、画像パッチとの明示的な位置対応を持たない「1次元トークナイザ」という新しい画像トークナイザのパラダイムを提唱しています。Vision Transformerをベースにしたアーキテクチャが採用され、情報を極めて小さく圧縮しながらも、高い再構成精度が得られています。このTiTokは画像生成においても高速かつ高品質な生成を実現できることがわかっています。

TiTokのアーキテクチャ
出典: https://arxiv.org/abs/2406.07550

One-D-Piece

こうした背景から、私たちは「圧縮率の高い画像トークナイザ」という領域に注目し、様々な研究を行っています。

私たちが提案した離散画像トークナイザである「One-D-Piece」はこの研究の成果の一つです。私たちはTiTokの「32トークン」という高い圧縮率が「高すぎる」という懸念を持ちました。圧縮率が高い場合、画像単位では十分な性能であっても、動画に対応させようとするとフレーム間の僅かな差異を上手く表現しきれない可能性があります。理想的には、「複雑な画像はたくさんのトークンを使い、単純な画像は少ないトークンで済ませる」ということを実現したいと考えられます。

そこで、TiTokに対し、「Tail Token Drop」と呼ぶシンプルな手法を組み込んだ新たな画像トークナイザを構築し、画像を可変長のトークン列にトークナイズすることを実現しました。SentencePieceやWordPieceのようなトークナイザの命名慣習に倣いつつ、1次元(One-Dimensional)トークナイザであることから、「One-D-Piece」と命名しました。

One-D-Pieceのアーキテクチャ

手法

「可変長圧縮」と呼べる画像圧縮手法は画像フォーマットの世界ではすでに十分に一般的です。よく知られるJPEGやJPEG 2000、WebPなど、多くのフォーマットが圧縮時に品質を制御するオプションを持っており、これを介して品質と圧縮率のトレードオフを調整することができます。一方、こうした手法は画像トークナイザの世界ではまだあまり導入されていません。

そこで注目したもう一つの技術は「Tail Drop」と呼ばれる既に提案されていた正則化手法です。オリジナルのアイデアでは、埋め込みベクトルの「末尾」ほど高いドロップアウト率を割り当てることにより、重要な情報をより先頭側に集める、ということを狙うものでした。これは表現学習におけるMatryoshka Representation Learningに類似した仕組みです。

https://ieeexplore.ieee.org/abstract/document/9174523

この「Tail Drop」をトークナイザに適用することによって、重要な情報が先頭に集まり、末尾は必要に応じて省略できる、「可変クオリティ」の画像トークナイザを実現することができるはずです。このアイデアをもとに「Tail Token Drop」というトークナイザ版のシンプルな正則化手法を考案しました。

実験

One-D-Pieceのモデルとして、S-256、B-256、L-256の3つのサイズを用意しました。これらはトークナイザのパラメタサイズに違いがあり、Sが最小、Lが最大のサイズです。一般に、より大きなモデルはより良い再構成の性能を達成することが知られています。
それぞれのモデルは最小1トークン、最大256トークンに画像をトークナイズすることができます。これらのモデルをImageNetを用いて学習し、評価します。学習はH100x8基で行い、直列換算で20日程度を学習に要しました。

画像トークナイザのおもな用途は画像生成であるため、基本的に再構成性能と下流タスク性能(特に、画像生成モデルの性能)で評価することが一般的です。この記事では再構成性能について主に紹介していきます。

まずは実際の振る舞いを示します。再構成についてはJPEGやJPEG 2000、WebPなどの画像フォーマットと比較しても圧倒的に良好な結果が得られました。これらの手法は可変品質の圧縮を可能にしていますが、画像トークナイザが活躍する高圧縮のドメインではあまり上手くいかないことがこの比較からわかります。また、TiTokは可変長圧縮に対応していないため、トークンを無理やり減らすと再構成が崩壊してしまいます。

画像トークナイザの再構成性能の評価指標として一般的なrFIDとPSNRを示します。特にrFIDは知覚的な品質を測る上で重視される指標ですが、One-D-Pieceで256トークン全てを用いる場合は近年の最先端のモデルと比較して同等以上の再構成性能を得られることがわかります。また、モデルサイズを大きくすることで性能が順調に上がることがわかります。

トークン数(または対応するbit数)を増やしながらrFIDを描画したのが以下の図です。

これに加え、様々な画像の再構成評価指標をトークン数を増やしながら評価したものが以下の図です。トークン数が長くなるにつれ、どのモデルでも性能が改善していくことがわかります。

One-D-Pieceは良好な再構成性能を示していますが、いくつかの弱点もあります。

まず、PSNRのような指標では圧倒的に従来の画像フォーマットが有利です。これはPSNRのような指標が人間の知覚よりもピクセルレベルの再構成精度を重視していることに由来するもので、One-D-Pieceは知覚的品質をより重視したことによってピクセルレベルの再構成品質を劣化させていると言えます。実はこうしたImage FidelityとImage Qualityの間のトレードオフは複数の文献で指摘されています。

また、同じトークン数、同じパラメータ数で比較した場合、固定長で学習されたTiTokにrFIDで劣ることもわかりました。こちらも可変長をサポートしたことによるトレードオフと考えることもできますが、類似の手法を導入した研究では性能劣化がほとんどないと報告するものもあるので、今後改善を目指したい点です。

分析

Tail Token Dropが画像の可変品質でのトークナイズを実現したことは、その挙動から明白です。しかし、学習されたトークンが実際にどのような性質を持っているのかははっきりわかりませんでした。

そこで、論文ではそれぞれのトークンの「寄与」を評価する分析を行いました。

分析の基本的なアイデアは、トークン列中のあるトークンt_iを別のトークンt_i'で置き換えたとき、画像がどれだけ変化するかを調べることです。もしそのトークンが重要な情報を担っていた場合、トークンを置き換えることによって画像が大きく変化するはずです。一方、大した情報を持っていないトークンであった場合、トークンを置き換えたとしても画像はほとんど変化しないことを期待できます。

そこでこのアイデアをもとに、評価用データに画像トークン列にランダムな置換を適用し画像の変化量の平均を取ることで、それぞれの位置のトークンの寄与を計りました。Tail Token Dropをかけたモデルでは「前方に重要な情報を持つトークンが集まる」ことを期待できます。

以下の図はこれを実際に可視化したものです。Tail Token Dropを適用したOne-D-Pieceでは左の方により強いピークが集まっている一方、このような仕組みのないTiTokではランダムな位置にピークが存在することが見て取れます。

また、トークン列にランダムな置換を適用した前後の画像の変化量を評価する上で、画像全体で平均を取る代わりに、位置情報を保ったまま寄与を可視化する、ということを行ったのが以下の図です。左上から右下に向けて順に見ていくと、最初の方のトークンは全体にぼんやりとした「大域的な情報を持つトークン」である一方、途中から特定の箇所を中心に広がる「局所的な情報を持つトークン」が増えていくことがわかります。このような傾向からも、Tail Token Dropが重要な情報、特に大域的な情報を先頭に集めてくれることがわかります。

このような結果を見ると、では「最初のトークン」は実際にはどんな情報に対応しているのかが気になってきます。そこで、先頭のトークンで画像を分類し、その一部を示したのが以下の画像です。どうやら「最初のトークン」は背景に強く対応しているようで、例えば209番トークンを最初に持つトークンは「白い背景の画像」ばかりでした。

「209番トークンが白い画像に対応する」ならば、背景が白くない画像であっても最初のトークンを209番に置き換えることで画像が編集できそうです。これを試してみたのが以下の画像で、最初の16トークンを用いて再構成する際に1番目のトークンを置き換えていくと、対応する背景色の画像を得られることがわかりました。残念ながらこの効果はトークン数を増やすと消えてしまうのですが、こうした挙動がみられるのは興味深いことです。

まとめ

この記事では、Turingが提案した画像トークナイザ「One-D-Piece」の開発背景と、その仕組み、さらに評価と分析について簡単に紹介しました。

今回の論文ではこのOne-D-Pieceというトークナイザの提案と評価、分析に焦点を絞りましたが、今後はこのような画像トークナイザを実際に世界モデルに適用していく予定です。

また、近年では画像と言語の混じった「ミックスモーダル」なデータを生成するInterleaved Text-Image Understanding and Generationの分野の研究が盛んに行われており、One Modelでこれを実現するGemini 2.0のNative Image Outputなど、実用水準の技術が実現されつつあります。中でもChameleonEmu3のようにテキストトークンと画像トークンの混ざった語彙空間でマルチモーダルな自己回帰モデリングを行うものは有望な方向性の1つです。将来的にはこのような用途で画像トークナイザが有用となることも期待できます。

最後に、Terraは画像を1フレームずつ生成する世界モデルですが、高性能な「動画トークナイザ」があればこれを「1秒ずつ生成」のような形に効率化できる可能性があります。Turingでは動画トークナイザにも注目しており、One-D-Pieceのような仕組みを持った動画トークナイザを実現するのも、今後の方向性の一つとなるでしょう。


Turingの基盤AIチームでは強力な機械学習モデルによって完全自動運転への道筋を確かなものにするための研究開発に取り組んでいます。この目標に取り組むMLエンジニア・リサーチャや、インターンを積極的に募集しています。

https://tur.ing/jobs

Tech Blog - Turing

Discussion