What's with ML software and pickles?

2023-11-08

MLソフトウェアはなぜpickleだらけなのか？

原文は Nelson Elhage により 2023-11-08 に公開されました。このブログを購読する

私は長年、機械学習とはまったく無縁のソフトウェアエンジニアとして過ごしてきたが、好奇心はあり、周辺的に関わることも時折あった。この間、私にとって繰り返し現れたテーマは、PythonのMLエコシステムでPythonのpickleが広く使われているのを目にするたびに覚える戦慄（正直に言えば軽蔑）だった。

重大なセキュリティ問題¹に加え、pickleをシリアライゼーションに使うと非常に脆くなり、コードを発展させたりライブラリやPythonのバージョンを上げたりする際にあらゆる悪夢を引き起こす。ソフトウェアエンジニアとしてのキャリアの中で、私は——良い近似として——pickleモジュールには正当なユースケースは存在しないと考えており、それは単なる魅力的な厄介者、より若く、より熱狂的で、より未熟だったエコシステムの過ちに過ぎないと信じるようになっていた。

しかしここ数年、私はMLの分野でプロとして働いており、エコシステムに対する新たな視点を得た。なぜMLソフトウェアが今のような形になっているのか、そしてpickleがユーザーのどんな問題を解決しているのかを、実践を通じて理解するようになった。今でもpickleは好きではないが、問題領域に対する理解と共感ははるかに深まり、もはやこの問題が「些細な」ものだとか、単なる無知や怠惰によるものだとは思っていない。この記事は、私が学んだことを少しでも伝え、二つの世界の隔たりを少しでも埋めようとする試みだ。

研究であって、ソフトウェア開発ではない

根底にある、私が最も伝えたい視点は、機械学習エコシステムのソフトウェアの多くは、何よりもまず研究者を支援するために設計されており、ソフトウェアエンジニアのためではないということだ。さらに、この文脈における研究は（少なくとも）いくつかの相互に関連した性質によって特徴づけられる（あるいは定義されるとさえ言える）：

研究の主要な成果物はソフトウェアではなく知識である。研究チームがソフトウェアを書くのは、ソフトウェアという道具やソリューション自体を持つためというよりも、研究上の問いに答え、自分自身やチーム、あるいは分野全体のドメインに対する理解を深めるためである。
研究者は、どのアイデアが成功するか、あるいは有用であるかを事前に知っていることはめったになく、だからこそ多くのアイデアを効率的に試すことを優先し、そのほとんどがうまくいかないことを前提としている。

上記の当然の帰結として、研究者（特にMLのような領域では！）は大量のコードを書くこともあるが、その大半がすぐに捨てられるのはまったく普通のことだ。失敗した実験はたいてい破棄され、成功したものでさえ論文を書き終えた後に捨てられることが多い。成功したアイデアを基に次の研究を進める場合でさえ、問題への理解が深まり、新たな問題や研究課題を念頭に置いて、コード自体は新しいプロジェクトのために書き直されるのが常である。

この環境におけるトレードオフやアフォーダンスは、典型的なソフトウェアエンジニアリングの文脈と比べると、（他のシリアライゼーション戦略と比べて）はるかにpickle寄りに傾いている：

ほとんどのコードは、同じ少数の人々（研究者とそのチーム）だけによって書かれ、単一の環境でのみ実行される。「開発」と「本番」の区別はほとんどないかまったくなく、外部のユーザーやデータソース、相互作用も存在しない。これらの要因が相まって、pickleによるシリアライゼーションの実践上のセキュリティリスクははるかに低くなる（ただしゼロではない！）
データ構造——あるいはコード断片さえ——を、スキーマやフォーマット、データの進化について立ち止まって考えることなく、そのままシリアライズできることには莫大な利点がある。重要なのは結果が出るまでの時間であり、スキーマやシリアライゼーションフォーマットについて考えたり、ディスク上の形式とメモリ上の形式の間で変換したりすることに費やす労力は、完全に本質的でない作業なのだ。
ほとんどのコードはかなり早い段階で捨てられるため、pickleの脆さや、コード変更やバージョンアップへの対応の難しさは、はるかに目立たなくなっている。

ある研究の情景

上記のトレードオフをもう少し具体的にするため、ML研究に取り組むある研究者の様子を様式化して描き、彼女がpickleに見出す具体的なユースケースと、私自身がML研究を行い、（はるかに経験豊富な！）ML研究者を支援してきた経験から見えるトレードオフを併せて示してみたい。

ある研究者が、新しい深層学習モデルのアーキテクチャ、あるいは既存のもののバリエーションを試しているとしよう。彼女のアーキテクチャには、数多くの設定項目やハイパーパラメータがある。層の数、各層の種類、さまざまなベクトルの次元数、活性化をどこでどのように正規化するか、どの非線形性を使うか、といったものだ。モデルの構成要素の多くはMLフレームワークが提供する標準的な層だが、研究者はそこに斬新なロジックをあちこちに差し込んでいくことになる。

私たちの研究者は、特定の具体的なモデル——これらの設定の特定の組み合わせ——を記述し、シリアライズして後で再読み込みできる方法を必要としている。これにはいくつかの関連する理由がある：

彼女はおそらく、ジョブを実行するためのGPUやその他のアクセラレータを備えた計算クラスタを利用できる。クラスタ上で動作するコードにモデルの記述を渡し、自分のモデルをクラスタ上で実行させる方法が必要だ。
それらのモデルが学習している間、ハードウェアの故障やジョブのプリエンプトに備えて、進捗のスナップショットを再読み込みして再開できる形で保存する必要がある。
モデルの学習が終われば、研究者はそれらを再び読み込み（最終的なスナップショットだけでなく、学習途中のチェックポイントも含めて）、評価や実験を実行したいと思うだろう。

ここではさまざまな解決策が考えられる！モデルのprotobufやJSONによる記述を書き、それを実際のモデルインスタンスに変換するコードを書くこともできる。モデルを「コマンドライン引数の配列」として表現し、それをパースしてモデルインスタンスにするという方法さえ考えられる。

しかし、これらすべての問題を解決し、しかも事前の追加作業がほぼ不要な解決策が一つある：

「特定のモデル」の定義は、単に「そのモデルを実装したオブジェクト」（おそらくtorch.nn.Moduleのサブクラス）である
研究者は、何らかのアドホックな方法で、実験を重ねる中で編集し発展させていく研究用スクリプトによって、これらのオブジェクトを構築する
そして彼女は、そのオブジェクト全体をpickleでシリアライズし、それを後の解析用やリモートマシンでの読み込み・実行用の保存バージョンとして使う。

このアプローチには面倒な手続きもいらず、研究コードや実装をどう構造化すべきかという特定の規律やパターンも不要で、「ただモデルを構築して実行する」ことと比べて認知的オーバーヘッドもほぼない。標準的なコンポーネントもカスタムなコンポーネントも同様に簡単に扱える。研究者はクラス階層の深いところでランダムな箇所に小さな調整を加えることも、モデルアーキテクチャ全体を置き換えることもでき、どちらの場合もシリアライゼーションの仕組みについて深く考えたり編集したりする必要はほとんどない。

この時点で、私たちの研究者はおそらく、さまざまなハイパーパラメータの直積グリッドをインスタンス化するスクリプトを書き、それらすべてをクラスタに投入して、学習が進む間に昼食を取ったり（あるいは週末休みに入ったり）するだろう。

これらのpickleファイルは確かに脆い！モデルの実装の詳細を丸ごとシリアライズするため、研究者が新しい設定やパラメータを追加したり、実験をリファクタリングしたりする際に、古い実験やモデルを簡単に壊してしまう。しかし繰り返しになるが、実験の大半は面白みのない結果に終わり、二度と見る必要はない。興味深い少数の結果については、各実験に関連するgitコミットを記録しておき、チェックアウトを戻せば十分だ。そして、もし特定のモデルをより永続的に保存する必要があるなら、どこかの__setstate__メソッドにちょっとした応急処置を加えることで大抵は修復できる。

とはいえ、実験が全体として成功し、この研究路線をさらに続け、より多くのバリエーションを試し、比較のためにモデルをさらに長く保存する必要が出てくると、この脆さは実際に痛みを伴うようになり、研究者はやがてリファクタリングを考えるかもしれない。単純なアプローチの一つは、モデルの「記述」と実装を分離することだ。モデルの設定・ハイパーパラメータをclass ModelConfig:にまとめ、このクラスこそがシリアライズ対象となり（依然としてpickleで）、研究者は与えられた設定から完全なモデルをインスタンス化する方法を知っているbuild(cfg: ModelConfig) -> Model関数を手作業で維持する。

研究者は今やこのクラスの前方・後方互換性だけを気にすればよく、ModelConfigからModelオブジェクトを構築するコードも併せてリファクタリングすることで、実装コードをかなり自由にリファクタリングできる。

ただし、私たちの研究者（そして／あるいは彼女のチーム）は依然として実験を続け、新しいことを試しており、反復と「実験までの時間」、そしておそらくそれ以上に「実験にかかる認知的コスト」がすべてを支配していることに留意すべきだ。だから、このクラスの多くのフィールドは素直な「普通の」整数や真偽値、あるいは列挙型になるだろうが、さらなる実験ごとに対応を追加しなければならないのは少々面倒なのだ。

そこで、研究者はおそらく、いくつかの「エスケープハッチ」を残しておく。ささいな例を挙げると、この分野では全結合層の後にテンソルに適用する非線形性、すなわち「活性化関数」を置き換えてみることが今でもそこそこ流行っている（私自身もここでは自白するとやってきた）。

では、活性化関数（それは文字通りTensor -> Tensor型のPython関数であることが多い）を設定にどう格納すればよいだろうか？

どこかのdict[str, Callable]に関数を登録して名前で保存することもできる……が、考えてみればPythonのモジュール名前空間自体がそういうdictのようなものではないか？だからpicklesを使っているなら、関数をそのまま設定に直接格納してしまえばよいのだ。

この気づきに勇気づけられ、研究者はさらに踏み込む。彼女はbuild_xxxx: Callable[[Config], Module | None]というオーバーライド全体さえ格納でき、これにより設定が設定の任意の部分を任意の方法でオーバーライドできるようになるのだ！

このような機能は、古い脆さの問題の大部分を呼び戻してしまう。しかし重要なことに、それによって研究チームは、より古く、より安定しており、「通常の」パラメータとして一級市民であるコアな設定と、（ほとんど）何も長続きせず、反復の速さがすべてを支配する実験の泡のような激しいchurnとの間で、バランスを取ることができるのだ。

スコープクリープ

私見では、上で述べてきたことのすべては、少々趣味が悪くリスクはあるものの、プロジェクトが真に信頼されたコンテキストで内部の研究用途にのみ実行される限りにおいては「基本的には問題ない」。

しかし、より大きな問題は、ひとたびこのようにpickleを日常的に使い始めると、あらゆるシリアライゼーションの問題にそれを使うのが習慣になってしまい、pickleがはるかに問題となるケースでさえそうしてしまうことだ。しかもそうするのは研究者だけではなく、このパターンはツールやエコシステム全体に漏れ出しがちなのである。

例えばPyTorchは、デフォルトのシリアライゼーションフォーマットの一部としてpickleファイルを使っている（torch.loadやtorch.save経由で）。しかし実際には、ほとんどの場合に本当に必要なのは、文字列とTensorのリストや辞書といった、ごく「プレーンな」データ構造だけである。PyTorchには、少数の組み込み型だけをサポートするカスタムのC++製pickleローダーさえあり、ローダーのweights_only引数で有効にできる。だが後方互換性のためにこのフラグはデフォルトでFalseになっており、（攻撃者がいない場合には）これを有効にしても壊れることしかない。そのためこのフラグはめったに使われず、すべてのPyTorchの重みファイルは、ユーザーが危険なインターフェースを使い、自らを脆弱なままにしておくことへのオープンな招待状となっている。

ではどうすべきか？

この記事を書く私の主な目的は、いわば「内側からの眺め」を浮き彫りにすること、すなわちなぜMLエコシステムの内部でこれほどpickleが使われているのか、そしてそれがユーザーのどんな問題を解決しているのかを示すことだ。とはいえ、機械学習の重要性が増し、さまざまな文脈でユーザーやデプロイが拡大するにつれ、安全でないpickleの使用によって引き起こされる問題の重要性と緊急性も増していくと私は予想しており、今後の道筋についていくつかの考察を提示したい。

より良いデータフォーマットとインターフェース

まず、私たち（ここで言う私たちとは、主にソフトウェアエンジニアリングのバックグラウンドやスキルを持つ者、あるいはMLの周辺で主にSWEとして働く者を指す）は、研究者により多く、より良い選択肢を提供するために、pickleに代わる高品質な代替手段への投資を続けるべきだと考える。うまくやれば、多くの場合、セキュリティ特性がより優れた上で、picklesよりもさらに良い体験を研究者に提供することさえできる。

ここで触れておかないのは不誠実だろうが、XLAを中心に構築されたGoogleの機械学習エコシステム——Tensorflowやjaxとそのさまざまなライブラリやフレームワーク——は、ここで述べてきた問題の多くを回避してきた。彼らはprotocol-bufferやnumpy-formatベースのシリアライゼーションフォーマットを使う傾向があり、pickleははるかに蔓延していない。一方で、私はそれらを、そのようなエコシステムが可能であることの存在証明のようなものと見ている。他方で、PyTorchがこれほど支配的になったのは、一般的に研究者により優れた使用体験、特に入門のしやすさや修正・探索の容易さを提供しているからだと私は考えており、Googleのライブラリを理想的な解決策として挙げることには躊躇いがある。

Huggingfaceのsafetensors libraryは、この分野における進歩の優れた例であり、テンソルデータを安全かつ高性能にシリアライズするための堅牢な基盤を提供している。私見では、それはやや低レベルなプリミティブと見るのが最適だ。dict[str, Tensor]という形のオブジェクトのシリアライズと復元をサポートするが、研究者はしばしばネストした辞書や、重みと一緒にメタデータを一つのバンドルに格納したいと思うものだ。safetensorsの上に構築される、より使いやすく研究者に優しいフォーマットやAPIの発展に私は期待している。

研究者はテンソルデータだけでなく、Pythonのクラスツリーもシリアライズしたいと思うことが多い。Pythonの型アノテーションと、その上に構築されたattrsやdataclassesのようなAPIの登場と近年の普及は、Pythonの型アノテーションと少量の追加メタデータ（おそらくcattrsを基盤として）によって駆動され、型アノテーションされたPythonクラスの［特定の］ツリーを多かれ少なかれ直接的にシリアライズ・デシリアライズできるライブラリの可能性をもたらすかもしれない。

`pickle`の制限

良くも悪くも、今日私たちは大量のpickleファイルと、pickleファイルを使うライブラリを抱えており、問題は依然として拡大している。代替手段の提供に取り組み、プロジェクトがそれらへ移行することを促す一方で、ここでは被害軽減あるいは緩和のための思索的なアイデアを一つ提示したい：

まず、Pythonがpickleモジュール全体を無効化するグローバルな「nopickle」モードをサポートすることが有用ではないかと考えている。そのようなフラグがあれば、少なくともPythonをデプロイする運用者は、実行時に信頼できないpickleを読み込むことで何らかのリスクに晒されているかどうかを容易に知ることができるだろう。

しかし、もしそれが制限的すぎるのであれば、pickleの使用を特定の方法で制限するオプションを追加することはどうだろうか？

Pickleマニフェスト

特に私が興味を惹かれるのは、pickleのすべての使用において、デコードのための有効なバイナリオブジェクトのリストを宣言することを必須とするアイデアで、おそらくsha256などの暗号学的ハッシュで名前付けされるものだ。pickle.loadの呼び出し元（あるいは何らかの新しいAPI）で期待されるオブジェクトを指定することと、帯域外の何らかのマニフェストファイルで指定することの両方をサポートすることが考えられる。マニフェストファイルはおそらく正当な（コールスタック、blob ID）のタプルを指定するだろう。このデータを指定しなかったり、許可されたハッシュと一致しなかったりするpickleの読み込みは、拒否されることになる。

このアプローチが興味深い理由は、MLソフトウェアをデプロイする人々が、一度だけ読み込まれるすべてのモデル重みファイルや設定ファイルのインベントリを収集し、そのセットを「凍結」することで、実行時に新たなpickleに遭遇しないことを保証できるかもしれないからだ。

そうなれば、pickleの使用を管理する問題は、大部分がサードパーティ依存関係を管理する問題に帰着する。これは確かに解決済みの問題ではないが、いずれにせよすでに抱えていた問題でもある。

具体的には、新しいサードパーティ製MLモデルをデプロイするためには、通常、モデル実装（Pythonモジュールやスクリプトの形で）とモデル重みの両方をインストールする必要がある。どちらも悪意があればセキュリティリスクとなるが、pickleによるリスクはある意味でより悪い。なぜなら、モデルコードの特定のバージョンをレビューしてピン留めしたとしても、pickleファイルはしばしば実行時にダウンロードされるため、依然としてシステム外部からのRCEに晒されることになるからだ。

もしそのコードが特定のpickleファイルだけを読み込むことを保証でき、その依存関係を初めて信頼すると決めた際に一度だけ監査すればよいのであれば、pickleのユースケースをサードパーティコードのインストール問題と「同じように見える」ものにできる。サードパーティコードでは、一度バージョンを取得してピン留めしてしまえば、少なくとも追加のランタイム攻撃面は存在しないのだから。

これらのアイデアを、ランタイム設定システムでさらに複雑にすることも考えられる。極限まで突き詰めれば、pickleの使用を制御するための次のようなかなり複雑なAPIを想像できるだろう：

with pickle.no_pickles() as token:
  pickle.loads(...) # NoPickleException

  with pickle.allow_pickles(token):
    pickle.loads(...)

  with pickle.allow_pickles(token, only_objects=[
    "sha256:some sha256",
  ]):
    pickle.loads(buf) # hashes buf, only proceeds if it matches

  # You must pass the same token returned by `no_pickles` to
  # `allow_pickle`; this prevents libraries from trivially blanket
  # calls to renable `pickle`
  with pickle.allow_pickles(bad_token):  # BadTokenException
    ...

このような仕組みが完璧になることはもちろんない。ライブラリがそれらを回避するハックを行うことは常に可能だ。極限を言えば、ライブラリはいつでもこれらの制限を一切持たないpickleの純粋なPython実装をベンダリングできる。しかし、それはライブラリが（意図的にせよ非意図的にせよ）バックドアを埋め込む数百万通りの方法のうちの一つに過ぎない。この提案の目的はサプライチェーンセキュリティを解決することではなく、pickleに特有のリスクを緩和し、前述の通り、リスクを実行時ではなくパッケージインストール時に移すことにある。標準的なセキュリティ機構の甚だしい迂回を実装するライブラリは、コミュニティによって気づかれ、信頼できないものとしてフラグ付けされることを期待したい²。

このような機能は、pickleモジュールをモンキーパッチするライブラリとしてツリー外でプロトタイプ化することも可能だが、標準化や普及という観点では、最終的にインタプリタ自体に組み込むことにも価値があるかもしれない。

10年以上前、私は（私の知る限り）pickleベースのシステムを攻撃するエクスプロイトについての最初期のpublic writeupsの一つを公開した。しかしこの問題は何年も前からよく知られており、今日ではpush-button exploitationのためのツールさえ存在する。↩︎
この文は、私自身が書いていてもナイーブに響く。しかし繰り返すが、そのような解決策が問題を既存のサプライチェーンセキュリティの問題へと大部分帰着させることは、良くも悪くも真実だと私は考えている。↩︎