Generatived（Beta）｜生成AIの最新ニュースとトレンドを提供

JAXは、AIイノベーターにとって強力なツールとして登場し、様々なアプリケーションにまたがる複雑なモデルのトレーニングを可能にしています。そのパフォーマンスとスケーラビリティは、高性能動画生成モデルでクリエイターエコノミーの分野で躍進を遂げているLightricksのような企業にとって特に有益です。Yoav HaCohen氏とYaki Bitterman氏が率いるLightricksのLTX-Videoチームは、以前のフレームワークでスケーリングに関する大きな課題に直面しましたが、JAXへの移行によって解決しました。

Lightricksは、生成AIによるコンテンツ制作の強化を目指していましたが、130億パラメータのモデル開発時に行き詰まりました。チームが以前使用していたフレームワークでは、必要なトレーニングステップ数とスケーラビリティを提供できなかったのです。JAXへの戦略的な移行は、即座に改善をもたらし、効果的なスケーリングと1日あたりのトレーニングステップ数の40%増加を実現しました。この変更により、モデル開発が加速されただけでなく、チームの生産性も倍増しました。

JAXの成功は、Flax、Optax、Orbaxといったツールを含む、シームレスに連携する包括的なエコシステムによるものでした。Lightricksはこれらのツールを活用して堅牢な新しいスタックを構築し、効率的な研究開発を実現しました。また、トレーニングにはGoogleのTPUコアを活用し、統合システムの設計を最大限に活用することで最適なパフォーマンスを実現しました。

LightricksのJAXに関する経験は、このフレームワークがエンジニアリングの課題を乗り越え、大幅なパフォーマンス向上を実現する上でチームを支援する能力を強く示しています。GoogleのJAXチームは、科学コンピューティングコミュニティの支援に引き続き尽力し、ユーザーが自身の経験を共有し、フレームワークの進化に貢献することを奨励しています。