IOAI2024 (At-Home Task) 日本代表ソリューション
今年の8月にブルガリアで開催されたInternational Olympiad in Artificial Intelligence (IOAI) に日本代表選手団の一人として参加してきました。
IOAIは国際科学オリンピックの一つで、AIに関する科学的技能を競い合う大会です。今年が第1回の開催で、34カ国のチームが参加しました。
チームメイトのnoboru君が素晴らしいIOAI参加記を書いてくれたので、ぜひご覧ください。
今回のIOAIには、モデルを訓練し、性能で競うScientific Roundと、AIツールを使った作品を作るPractical Roundがありました。Scientific Roundは、Kaggleのような形式と考えていただいて大丈夫です。どちらのラウンドもチーム戦となっていて、日本からは4人の1チームで参加しました。
この記事では、Scientific Roundの特にAt-Homeタスクのソリューションを共有します。私はNLPとCVを担当したのですが、MLをnoboru大先生に任せてしまったので、ここではNLPとCVにのみ限ります。
Scientific Roundのスケジュールは次のようになっていました。3つのタスクはML (Tabular), NLP, CVが与えられました。
- 6月 - 8月: At-Home Task
- 3つのタスクにじっくり取り組む
- 8月@Bulgaria: On-Site Task
- At-Home Taskをベースとしたタスクに8時間で取り組む
Kaggleとの相違点を挙げてみます。
- 問題はサンプルの.ipynbで与えられる
- Leaderboardはなし
- モデルのweightと評価の.ipynbを提出
- 全て見えないテストセットで評価
- 質問は国のteamleader (選手とは別)経由でslackでjuryに投げる
- Google Accountが与えられ、T4が数百時間？回せる程度のクレジットのみ使って問題に取り組む。外部リソースは禁止されている。
問題文はこのページから入手することができます。
NLP
NLPは暗号化された言語で書かれた文章の5値分類タスクでした。
データセットは2種類あり、1つはラベル付きデータセットで、もう1つはラベルなしのコーパスでした。
さらに、次の制約が設けられていました。
- mBERTをベースとしたモデルのみ使用可能
- 訓練はL4 GPUで8時間以内
- 推論は500個のサンプルに対して5分以内
文章は全て次のようにブラーフミー文字とデーヴァナーガリー文字を組み合わされていました。問題文にはこれを復号することはできないし、やる意味がないと書いてありました。おそらく、既に存在している言語の文字をランダムにインド系の文字に置き換えただけなのだと思います。
ラベル付きデータは1.5k行程度あるのに対し、ラベルなしのデータは61k近くありました。
| 文 | ラベル |
|---|---|
| झ𑁣झच𑀪𑀢𑀟 𑀣च 𑀠न𑀞𑁦 ण𑀢 𑀟च 𑀫चझ𑁣 𑀠च𑀟 𑀲𑁦पन𑀪 च ब𑁣𑀠ढ𑁦 𑀣च ढचनत𑀫𑀢 ष ब𑀱च𑀠𑀟च 𑀢𑀟न𑀱च णच𑀫चणच | 0 |
| णच𑀣𑀣च 𑀞णच𑀟𑀣च 𑀞𑁦 ढच𑀪च𑀤च𑀟च च 𑀣न𑀟𑀢णच |
4 |
| ढचढच𑀟ब𑀢𑀣च चल𑀢णन𑀕 𑀙णच𑀟 𑀟च𑀘𑁦𑀪𑀢णच 𑀟च 𑀞𑁦𑀱च𑀪 𑀳𑀫𑁦𑀞च𑀪न 𑀭𑁢 𑀟च 𑀠नल𑀞𑀢𑀟 ध𑀣ध |
3 |
私たちのアプローチは以下の通りです。
- tokenizerの再訓練
- コーパスを使ったmBERTの継続事前学習 (MLM)
- 訓練時間短縮の工夫
はじめに、mBERTでモデルを訓練した際、なかなか性能が向上せずデバッグしたところ、入力テキストの多くがUnknownラベルになっていることに気付きました。
tokenizerの再訓練が必要なことがわかったので、この記事を参考に、再訓練しました。sentencepieceでいいのか迷いましたが、テキストの中にスペースがあることから、sentencepieceでやることにしました。
また、MLMの継続事前学習については、Hugging Faceのサンプル実装を参考に実装しました。Data Collatorを定義し、 BertForMaskedLM を呼び出すだけで訓練ができました。
最終的なパラメータです。
- Pre Training (MLMの継続事前学習)
- 2 epochs
- Batch Size: 32
- Max Length: 256
- optimizer: AdamW
- Learning Rate: 5e-5
- scheduler: Linear
- 半精度
- Fine Tuning
- 30 epochs
- Batch Size: 32
- Max Length: 256
- optimizer AdamW
- Learning Rate: 5e-5
- scheduler: Linear
- 単精度
うまくいかなかったこと
- Transliteration (翻字): 与えられてる文字をラテン文字に置き換える
- Pseudo Labeling: コーパスにラベルを割り振ると9割のテキストが単一のラベルに割り当てられ、不均衡になり、精度が上がらない
- 切り詰め: 実装が間に合わなかった
余談、提出後の気付きによるスコアアップ
日本チームが提出したAt-Home Taskのソリューションは以上なのですが、提出したあとにスコアアップしたので残しておきます。
きっかけは、バングラデシュのチームがHugging Faceにデータセットや重みを公開していたのを見つけたことでした。彼らはtransformersの Trainer メソッドを使っていたので、F1スコアを含むログもアップロードされていました。確認すると、我々よりも高いスコアを出していたので、気になってモデルを調べていました。
彼らの仕事を整理すると、
- モデルサイズは我々の20倍くらいある
- 増加分は全てembedding層だった
- embeddingの次元を増やした + fasttextを使用
- tokenizerも奇妙
- eval f1は0.97近くを主張している
でした。
また、彼らの vocab.txt は完全にwordレベルのトークナイズが行われていました。
チームのありさ氏がfasttextのembeddingをBERTのembeddingにすり替えているのではという指摘があり、2人で調査したところ、その通りでした。彼らのスコアの本質はfasttextにありました。
embeddingの次元は (vocab_size, embeding_dim) で与えられます。
我々は、 (33297, 768) でしたが、彼らは (548133, 768) でした。また、fasttextで得られる埋め込みベクトルはデフォルトで300ですが、可変であることも確認しました。
彼らのvocab_sizeが異常にでかいのは置いておいて、fasttextを使ってBERTのembedding vectorを置き換えることに成功しました。
fasttextを使ってembedding vectorを置き換えて、事前学習した後に、fine-tuningをしました。下の図はfine-tuningのときのepochごとのf1スコアを示しています。
緑はfasttextのembeeding vectorをBERTに置き換えたもので、下の青が最終提出したモデルです。最初から比較的高いスコアを獲得し、全てのepochで我々の最終提出を上回っていました。
なぜembedding vectorの置き換えが有効だったのでしょうか。それは、トークナイズとembeding vectorの対応のミスマッチを解消するからだと考えています。
我々はmBERTの学習済みモデルを使用したため、本来のトークナイズで得られるtokenとそのtokenのidのembedding vectorは対応しているはずです。しかし、tokenizerのみを再訓練してしまったため、embedding vectorとミスマッチが置きていました。
CV
CVのテーマは拡散モデルでした。既存のtext2imgモデルである、 lambdalabs/miniSD-diffusers を再訓練して "giraffe" と "zebra" の意味を入れ替えろというタスクでした。
勘の良い方なら、embeddingやtokenizerを調整するだけで解決すると考えるかもしれませんが、これらの操作は禁止されており、純粋な訓練で対応する必要がありました。
サンプルのノートブックはCOCO2014のデータセットを使って、captionを入れ替えて再訓練するチュートリアルが与えられていました。
このタスクの面白いところは評価指標で、プロンプトに含まれているオブジェクトのみが生成された画像に含まれているかどうかを物体検知モデルで判定し、accuracyをスコアとしています。
例えば、 A giraffe is running on the grass というプロンプトを入力したらシマウマが出てくることが期待されます。
検出結果とスコアの対応は次のようになっており、オブジェクト以外の検出がされた場合にはスコアは0になってしまいます。
- (シマウマ): 1
- (シマウマ以外): 0
- (シマウマ, それ以外): 0
なお、この評価指標では画像のクオリティは全く考慮されていないことに注意してください。
我々のアプローチは次のとおりです。
- SDXL-turboを使ったデータセット拡張
- 評価指標ハック (オリジナルのYOLO Lossの導入)
- 信頼度の高いテストセット生成
以下はCOCO2014を使って再訓練を施したときの生成結果のサンプルなのですが、全体的にクオリティが低いことが伺えます。
COCO2014は最近の拡散モデル向けのデータセットと比較するとcaptionの質があまり高くないので、代替案を考えていました。
次のようなことを考慮した結果、人工データセットを構築することになりました。
- 訓練時間に制約がある
- 生成のクオリティは重要ではない
- アンリアリスティックなシチュエーションの画像も生成したい
人工データセット構築に使うために様々なtext2imgモデルを試しましたが、最終的にはSDXL-turboを選択しました。1stepでプロンプトに沿ったクオリティの高い画像を生成してくれました。
アンリアリスティックな画像に対しても高いクオリティで生成してくれます。L4 GPUで0.35秒 / 1枚の速さで生成してくれました。
生成に使うプロンプトは、YOLOのクラスをオブジェクトを含むようにしてChatGPTに生成してもらいました。合計で3000件の画像とキャプションのペアができました。
プロンプトのサンプルです。
"A baseball bat with a telescope.",
"A baseball bat at a science fair.",
"A baseball bat in a hospital.",
"A baseball bat in a firefighter's hands.",
"A baseball bat at a police station.",
"A baseball bat at a city hall.",
"A baseball bat in a flower shop.",
"A baseball bat in a bakery.",
"A baseball bat in a barber shop.",
"A baseball bat in a pet store.",
"A baseball bat at an amusement park.",
"A baseball bat with a roller coaster.",
"A baseball bat at a Halloween party.",
"A baseball bat in a haunted house.",
"A baseball bat at a holiday parade.",
"A baseball bat with Christmas lights.",
"A baseball bat at a New Year's Eve party.",
次に、評価指標のハックについて紹介します。このタスクでは前述した通り、生成された画像に指定されたオブジェクトが含まれているかどうかでスコアリングされます。
我々がとったアプローチは、訓練時に指定されたオブジェクトが含まれるように促す補助ロスを追加することでした。YOLOを使っているのでYOLO Lossと名付けています。
拡散モデルはt=0からt=Tまでノイズをステップバイステップでデノイズします。訓練時には、nステップ先 (
YOLO Lossの実装
# 以上省略
model_pred = unet(noisy_latents, timesteps, encoder_hidden_states, return_dict=False)[0]
# 物体検知モデルの定義
yolo_model = YolosForObjectDetection.from_pretrained('hustvl/yolos-tiny')
yolo_model.requires_grad_(False)
# デノイズするために必要なパラメータの計算 (下記リンクに従って)
# https://github.com/huggingface/diffusers/blob/v0.29.2/src/diffusers/schedulers/scheduling_ddpm.py#L499-L523
alphas_cumprod = noise_scheduler.alphas_cumprod
sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.view(-1, 1, 1, 1)
sqrt_alpha_prod = sqrt_alpha_prod.view(-1, 1, 1, 1)
# latentのデノイズ
denoised_latents = (noisy_latents - model_pred * sqrt_one_minus_alpha_prod) / sqrt_alpha_prod
# latentを画像空間に戻す
denoised_sample = vae.decode(denoised_latents.bfloat16() / vae.config.scaling_factor).sample
# YOLOで物体を検知する
outputs = yolo_model(denoised_sample)
logits = outputs.logits
prob = nn.functional.softmax(logits, -1)
# evaluationに合わせてスコアを計算する
batch_logits_list = []
for i in range(len(labels)):
list_logits = []
for j in range(len(labels[i])):
# Check if the label is within the specified range and the score exceeds the threshold
if (labels[i, j] in LABELS_RANGE) and (scores[i, j] > CFG.threshold):
list_logits.append(logits[i, j, :-1])
# If no logits were added to the list, append a zero matrix
if len(list_logits) == 0:
batch_logits_list.append(torch.zeros_like(logits[i, 0, :-1]))
else:
# Otherwise, sum the logits
batch_logits_list.append(torch.stack(list_logits).sum(dim=0))
batch_logits = torch.stack(batch_logits_list)
yolo_loss = ce(batch_logits, batch["labels"].to(device))
本来の拡散モデルのLossに補助する形で、YOLO Lossを導入して訓練したモデルの推論サンプルは次のようになっています。
同じ動物の重複は考慮されないため、際限なく増えています。
YOLO Lossの係数を調整すると、動物の密度も変わることがわかりました。
ここまでは、モデルの改良でしたが、実際の評価に使われるプロンプトは公開されていませんでした。公開されているのはサンプルのEvaluation Notebookのみですが、そこには8つのプロンプトしか公開されていないので、自前でValidationを組む必要がありました。
prompts = [
["A curious zebra standing tall in a lush African savanna at sunrise, with acacia tree,s in the background.", "giraffe"],
["Next to a medieval castle, a regal zebra observes the knights and a drawbridge.", "giraffe"],
["Wearing a scarf, a fashionable giraffe strolls through a bustling city street with skyscrapers.", "zebra"],
["Running along a sandy beach, a playful giraffe enjoys the palm trees, ocean waves, and a bright sunset.", "zebra"],
["By a serene lakeside, a relaxed bear drinks water with mountains and a clear blue sky in the background.", "bear"],
["In a snowy forest, a cozy bear stands under snow-covered trees, enjoying the gentle snowfall.", "bear"],
["Partially hidden in a dense tropical rainforest, an adventurous sheep peeks through leafy plants.", "sheep"],
["A sleek sheep with modern accessories navigates a futuristic city with flying cars and neon lights.", "sheep"]
]
これらのプロンプトを観察し、我々は複雑で現実/非現実のプロンプトを用意することに決めました。我々は100以上の様々なプロンプトを作り、疑似評価セットを作成してモデルを提出しました。その評価セットでは正解率が0.94になるまでスコアを上げることができました。
しかしながら、最終的なスコアは正確な記憶ではないんですが、中の下?程度でした。後から評価セットを覗いたところ、プロンプトは数ワードが多いとのことでした。つまり、長くて複雑なプロンプトに過剰適合させてしまった可能性が高いです。
YOLO Lossを使って評価指標ハックしようとしたのに、肝心の評価セットと公開されているデータとがかなり異なるのはどうなのかなと思います >.<
チームの動き方
Kaggleでは、一人ひとりが別々のアプローチでモデルを組んで最終的にアンサンブルするというやり方を多く経験してきましたが、IOAIに関してはモデルの制限があったのでアンサンブルはまったく使えませんでした。したがって、複数人が似た実験を回すことになります。その上で、チーム連携を強くするために次のことを意識しました。
- 最初に実験ベースラインの作成
- WandBの導入
- 1実験1スクリプトの徹底
実験のベースラインはY.NAKAMAさんのコードを意識しました。
- クラスでCFG管理
- colabで動かすので、yamlや環境変数は使えない
- 生PyTorchを使う
- デバッグしやすくするために、transformersのTrainerなどの便利な訓練ツールに依存しない
- 構造をシンプルにする
- 初期化 -> 訓練 -> 評価の流れを明確にする
以上に加えて取り組んだこともあります。
- WandBを使った実験管理
- 重みのHugging Face Hubアップロード
はじめのうちは、ベースラインを構築して複数人が実験しても条件が変わらないように意識しました。
WandBを使うことで、実験結果の解析に捗るのはもちろんなんですが、デバッグサイクルの高速化にも寄与しました。NLPの事前学習は4時間-8時間かかるのですが、実験を終わるのを待たずともLossの動きで実装のバグに気付くことができました。CVに関しても、Lossを見ただけではわからないものも、適当なstepごとに画像を生成してWandBに上げるようにすることで、定性的な妥当性評価もできるようになりました。
実験は、1実験1スクリプトを採用し、原則実験したままに残しておくようにしました。
colabは簡単に書き換えられてしまうので、GitHubに共有することによってコードを残すようにしていました。
最後にコミュニケーションについてですが、Discordの上で行いました。必要に応じてインスタンスな会議も開きました。
-
2024_athome_general: at homeタスク全般に関して -
2024_ml, nlp, cv: タスクごとの議論 -
2024_diary: 日報
おわりに
2ヶ月間、全力投球したコンペでした。。勝てませんでしたが、自分の中ではやれることはやりきったので、まだまだ実力が足りないようです。コンペ期間通して、チームメンバーも本当に実装力が強くて非常に支えになりました。
これからのToDoなのですが、私が言及しなかったon-siteタスクについてや、best solutionの解説についてもいずれできたらと思っています。
IOAIの実験に使ったリポジトリは公開したので、もし実装が気になる方がいたらこちらからアクセスしてみてください。
来年のIOAI2025は北京で開催されるそうです。来年も日本から代表団を送るそうなので、気になっている高校2年生以下の皆さんはIOAI-JAPANのwebsiteやTwitterをチェックしてぜひ参加してみてください！
IOAIに関して質問がある方は遠慮なく私にも聞いて下さい！
Discussion