HunyuanVideoのDualStreamBlockについて

1.HunyuanVideoの構造

HunyuanVideoにおけるDualStreamBlockという構成がある。
これに既視感があったのでどこで見たかを思い出してみる。

ここでパラメータの1/4はTextToken(MLLM)のDualStreamBlock、同じく1/4はVideoToken(3D VAE)のDualStreamBlock、残りの1/2はSingleStreamBlock(DualStreamBlockの層深さ2倍)であろう。

論文の最後の方でFluxに近いdual flow attention blockを用いていると言及されている。

2.他モデルとの比較（Flux.1, SD3, VLM）

DoubleStreamBlock(Flux.1)

画像生成AIであるFlux.1にはDoubleStreamBlockとSingleStreamBlockという構成がある。
Flux.1のダイアグラムではDoubleStreamBlock19層とSingleStreamBlock38層の構成らしい。
HunyuanVideoはDoubleStreamBlock20層とSingleStreamBlock40層の構成で、パラメータ数がFlux.1が12BでHunyuanVideoが13BなのでDiT本体の構造はおそらく近い。

またFluxの音楽版(FluxMusic)の論文の図に以下がある。
上述のダイアグラムより図としては分かりやすい。

StableDiffusion3(MM-DiT Block)

また、元をたどればStableDiffusion3におけるMM-DiT Blockである。

mmproj(VLM)

MM-DiT BlockのMMは(multi-modal)の略である。これはマルチモーダルLLM、MLLM(VLM)において聞き覚えがある。この時のmmprojは単なる2層のMLP層だったがこれをTransformerに変えればHunyuanvideoのDualStreamBlock(EncoderとLLMを接続する)に近くなるのではないだろうか。

イメージ図

DoubleStreamBlockと、VLMにおけるmmprojとの類似性を考えた図を以下に示す。

また、Tフレームの画像をCLIP ImageEncoderに通すと(W/14 × H/14, T, 1024)になる。一方、3D VAEに通すと(W/8 × H/8, (T+3)/4, 16)となり、同じEncoderでも圧縮次元が異なる。

この図で見るとDoubleStreamBlockがImageEncoderとLLM(Transformer)を仲介するmmprojのように3D VAE EncoderとDiT(SingleStreamBlock)の仲介するように見える。

3.TripleStreamBlockとMoE

ここからはポエムである。

TripleStreamBlock

前述の図を考えるとDoubleStreamBlockではなくTripleStreamBlockとすればLlavaのImageTokenか、もしくはCLIPTextEncoderのTextTokenをVLMのImageTokenのように動画生成モデルにそのまま取り込めるのではなかろうか。

またはImageTokenStreamBlock、VideoTokenStreamBlock、TextTokenStreamBlockの三つを構成し、Video学習の時はImageTokenStreamBlockを固定化し、Image学習の時はVideoTokenStreamBlockを固定化すれば、初期の動画学習におけるDualStreamBlock上の破滅的忘却を回避できるかもしれない。

HunyuanVideoでは音楽生成モデルにおいてTripleStreamBlockは提案されている。

ImageとTextの各Self-Attentionを$I^2,T^2$
ImageとTextのCross-Attentionを$IT$とすれば
ImageとTextのTokenを結合した場合のSelf-Attentionは$(I+T)^2$であり、各Self-Attention項を差し引くことで$(I+T)^2 - I^2 - T^2$となる。この結果は、ImageとText間のCross-Attentionを結合TokenのSelf-Attentionとして表現できる可能性を示唆している。

MoE(Mixture-of-Experts)

LLMにおいては大型なLLMを学習するより小型のLLMを事前学習してその重みをコピー複製して、GateによってToken毎に異なるLLMに割り振りファインチューンすることによって少数Tokenにモデルを最適化するMoEが近年脚光を浴びている。

DeepSeek-V3は37BのLLMを複製して671Bモデルを構成している。
動画生成モデルでもMoEによってDiTを複製してtokenごとに振り分け、描写タスクごとにFineTuneした方が最初から頑張って大型モデルを構築するより低コストで学習可能なのだろうか。例えば12BのDiTをはじめから学習するより2B DiTモデルを作って学習し、これをコピーして6個に複製してGateによりTokenを6通りに振り分けて各DiTにおいてFineTuneする方が12Bを上回る性能を示すかはともかく学習コストは低いと考えられる。(モデルパラメータ数が6倍になれば通常の学習コストは36倍になるが、MoEでは学習コストが増加するのはFine-tune部分のみで、それも最大6倍程度に抑えられる)

また、SDXLのrefinerやeDiff-Iのexpert denoisersとかがnoise levels(推論step)に応じてモデルを変えるということはやっていたがTokenに応じてモデルを変えるというのは知らない。この異なるImageTokenとTextTokenの最適化こそがDualStreamBlockに相当し、また単純なSingleStreamBlockのみのDiT学習は高コスト(出来なくはないが物理パワーで殴るようなもの)であり、DualStreamBlockのMM-DiT学習はMoEのように低コストで済むものなのだろうか。

下記図はStableDiffusion3の論文より

4.画像生成モデル重みは使い回せないのか？

HunyuanVideoは既存の画像生成AI(特にSDXLの派生)にくらべあまり日本のキャラクターの生成は得意ではないようだった。HunyuanVideoのモデルサイズは13Bなのに0.9BのStableDiffusion1.5と同程度でしかない(もしかしたら学習キャプション上の問題かもしれないが)。おそらくこれはHunyuanVideoがまず一から画像生成モデルを学習して、次に動画生成モデルを学習していることによる。

VLMにおいてはTextDecoder(LLM)は既存のLLM重みをそのまま使いまわしている。
ということは既存の画像生成AI(SD3やFlux.1)におけるSingleStreamBlockを重み固定にしてDualStreamBlock(DoubleStreamBlock)の浅い層のみを学習させ、最後にSingleStreamBlockとDualStreamBlockをファインチューンするのがVLMにならう学習の手順である。

従ってMM-DiT型な画像生成AIにおいてEncoderを2D VAEを3D VAEに差し替え、T5-XXLをllava-llama3に差し替える。次に3D VAE Encoder、llava-llama3とDiTの間の浅い層のDualStreamBlockと3D VAE DecoderとDiTの間の最終層付近のSingleStreamBlockのみを学習すれば3D VAE、llava-llama3をEncoderとし、3D VAEをDecoderとした画像生成AIモデルを比較的低コストで作成できよう。
その後にこのDiTを動画でDual,Single-StreamBlock学習すれば動画生成AIとなろう。

パラメータ数はSD3mediumが2B、SD3largeが8B、Flux.1devが12Bである。

こちらの方がFlux.1devの大部分の重みを凍結したまま学習出来る分コスト的に有利なのではと思った。

Token refiner

再度Hunyuanvideoの構成を注意深く見てみるとToken refinerなるものがMLLMとDiTの間にある。

これの詳細は参照論文に書かれており、ざっと見た感じ通常Llama3の特徴量をDiTに使っても何故か良くなく、位置バイアスの偏りを無くす必要があるらしい。

もし、Token refinerがLlama3出力をT5 likeにするなら、同様にあるEncoderを別のEncoderに差し替える際のToken refinerを下記の図で示す赤ブロックの様な水色の矢印間が一致するようにCLIPTextEncoder(とTimeEnbedding)をCross-AttentionとするSingleStreamBlockを各自学習して、これを前述したFlux.1のDiTの先に三つのToken refinerを結合するのも同じである。

5.学習コストの考察

仮に既存の画像生成モデルから動画生成モデルの学習を始められた場合、どの程度のコストが削減可能かを考えてみる。

HunyuanVideoの論文によれば13BモデルのT2I(画像生成)学習に$5.8×10^7 Peta(10^{15}) FLOPs$、T2V(動画生成)の学習に$7.0×10^7 Peta(10^{15}) FLOPS$という時間が読み取れる。
これは92M～6.6Bのモデルのスケーリング則から13Bのモデルサイズ時の学習時間を見積もった値であろう。

一方、StableDiffusion3の論文によれば8BモデルのT2I学習に$5.0×10^{22} FLOPs$とあるので学習FLOPsでいえば同じ程度である。Flux.1の学習コストについては情報はない。

H100のFP16性能は2ペタFLOPs(2000TFLOPs)らしいので$5.8×10^7/2/3600/24=335.6$日から
H100単体だと11ヵ月、64個で5日程度の学習時間と思われる。
awsだとH100×8が7000円/hくらいだからHunyuanVideoのT2Iの学習費用は概算で800万円くらい、T2Vの学習費用は概算で1000万円くらいとなる。

StableDiffusion2のフルスクラッチの学習コストが＄50k(800万)なので13Bのモデルサイズの割には低コストのように思う。

動画生成モデルが既存画像生成モデル重みからの学習が可能ならば1800万円の内、800万円を削減出来よう。

まとめ

HunyuanVideoのDualStreamBlockについて考え、これがFlux.1やStableDiffusion3の構造に近く、もっと言えばVLMにおけるImageEncoderとLLMを接続するmmprojに近いのではないかという考察を行った。
VLMが既存LLMの重みを活用できるのと同様に、動画生成モデルもDiTとEncoderの接続部を学習することで、既存の画像生成モデル（FluxやSD3）の重みを活用できるのではないか。