Databricks Feature Store — データ/MLOps プラットフォームの協調設計による特徴量ストア

Today, we announced the launch of the Databricks Feature Store, the first of its kind that has been co-designed with Delta Lake and MLflow to accelerate ML deployments. It inherits all of the benefits from Delta Lake, most importantly: data stored in an open format, built-in versioning and automated lineage tracking to facilitate feature discovery. By packaging up feature information with the MLflow model format, it provides lineage information from features to models, which facilitates end-to-end governance and model retraining when data changes. At model deployment, the models look up features from the Feature Store directly, significantly simplifying the process of deploying new models and features.

AI における「データ」という課題

特徴量エンジニアリングとは、データを機械学習モデルで処理できる特徴量に変換する作業のことです。未加工のデータ(トランザクションログ、クリック履歴、画像、テキストなど)をそのまま機械学習(ML)に使用することはできません。そのため、データエンジニア、データサイエンティスト、ML エンジニアは、未加工のデータから機械学習モデルで使用できる最終的な「特徴量」に変換するために多大な時間を費やしています。この「特徴量エンジニアリング」は、データの集計(例:特定の時間ウィンドウにおけるユーザーの購入件数)から ML アルゴリズムの複雑な結果(例:word embeddings)まであらゆるものが含まれます。

このようなデータ変換と ML アルゴリズムの相互依存性は、ML モデルの開発、デプロイに大きな課題を引き起こします。

  • オン/オフラインの違い:多くの意味のある ML ユースケースのいくつかにおいては、モデルは低いレイテンシでオンラインにデプロイされる必要があります(Web ページが読み込まれた際に数十ミリ秒で実行されるレコメンデーションモデルを考えてみてください)。トレーニング(オフライン)における特徴量計算に使用される変換処理は、モデルのデプロイメントの際(オンライン)には、低いレイテンシで処理される必要があります。このため、チームは特徴量計算を何度も実装しなくてはならず、モデルの品質に重大な影響をもたらす些細な違い(オン/オフラインの違い)を持ち込んでしまいます。
  • 再利用性と発見可能性:多くの場合、特徴量を探し出すのが容易ではないため、何度も再実装しなくてはならず、仮に発見できたとしても再利用しやすい形で管理されていません。この問題に対する素朴な解決策は、特徴量の名前に基づく検索機能の提供ですが、この場合、データサイエンティストは他の人が特徴量にどのような名前をつけたのかを推測しなくてはなりません。さらには、どの特徴量がどこで使われたのかを知るすべがなく、特徴量テーブルの更新・削除を決断するのが困難になります。

Databricks Feature Store(特徴量ストア)の独自のアプローチで AI におけるデータの問題を解決

The Databricks Feature Store is the first of its kind that is co-designed with a data and MLOps platform. Tight integration with the popular open source frameworks Delta Lake and MLflow guarantees that data stored in the Feature Store is open, and that models trained with any ML framework can benefit from the integration of the Feature Store with the MLflow model format. As a result, the Feature Store provides several unique differentiators that help data teams accelerate their ML efforts:

  • ネイティブモデルパッケージングによるオン/オフラインの違いの排除:
    MLflow との連携によって、Feature Store によるモデルのアーティファクトと特徴量検索ロジックのパッケージングを可能にします。Feature Store のデータを用いてトレーニングされた MLflow モデルがデプロイされる際、モデル自身が適切なオンラインストアから特徴量を検索します。これは最初から Feature Store が存在するかどうかを気にせずに、クライアントがモデルを呼び出せるということを意味します。結果として、クライアントの複雑性は低下し、モデルを呼び出すクライアントを変更しなくても、特徴量を更新できることになります。
  • 自動化されたデータリネージュ(data lineage)追跡による再利用性、発見可能性の実現:
    データネイティブな環境における特徴量計算によって、Databricks Feature Store が特徴量計算に用いられたデータソース、使用されたコードバージョンを自動で追跡できるようになります。これにより、データリネージュベースでの検索が可能となります。データサイエンティストは未加工のデータを用いることで、同じデータから計算された全ての特徴量を見つけることができます。加えて、MLflow のモデルフォーマットとの連携により、特徴量からモデルに至る下流のリネージュを提供します。Feature Store は正確にどのモデルと特定の特徴量を使用したエンドポイントなのかを理解しているので、エンドツーエンドのデータリネージュを実現し、特徴量テーブルに対する更新、削除を安全に意思決定できます。

ディスカバリー、コラボレーション、ガバナンスの中央リポジトリとしての Feature Store UI

Feature Store は企業における全ての特徴量の中央管理リポジトリとなります。全ての特徴量に対して検索可能なレコード、定義・計算ロジック、データソース、作成者、特徴量の利用者を参照できます。UI を用いることで、データサイエンティストは以下のことが可能となります。

  • 特徴量テーブル名、特徴量、データソースによる特徴量の検索
  • 特徴量テーブルから特徴量、接続されたオンラインストアへのナビゲーション
  • 特徴量テーブル作成に用いられたデータソースの特定
  • 特定の特徴量を使用する利用者(モデル、エンドポイント、Notebook、ジョブ)の特定
  • 特徴量テーブルのメタデータに対するアクセス権管理


Feature Store 検索は、特徴量名 'customer_id'、'feature_pipeline' を持つ特徴量テーブル、名前に 'raw_data' ソースを使用。

Databricks Feature Store は、他の Databricks のコンポーネントと完全に統合されています。このネイティブな統合によって、特徴量計算に用いられたデータ、特徴量をリクエストした全ての利用者、Feature Store の特徴量を利用してトレーニングされた Databricks モデルレジストリ上のモデルの全てのデータリネージュを可能にします。

特徴量テーブル user_features.behavior のデータソース、作成元のノートブック、モデル、利用側のモデル、エンドポイント、ノートブックを表示。
Feature table ‘user_features.behavior’, shows lineage to data sources and producers notebook; models; as well as consumer models, endpoints, and notebooks

高スループットでオフライン、低レイテンシでオンラインでの機能への一貫したアクセス

Feature Store は特徴量にアクセスするアプリケーションに対するさまざまなオフライン、オンラインの特徴量プロバイダをサポートしています。特徴量は 2 つのモードで提供されます。バッチレイヤーは、ML モデルのトレーニングとバッチ推論において特徴量を高いスループットで提供します。オンラインプロバイダはオンラインでのモデルサービングで同じ特徴量を低いレイテンシで提供します。オンラインプロバイダはプラグイン可能な共通的な抽象化で定義され、さまざまなオンラインストアをサポートし、オフラインからオンラインに特徴量を公開し、特徴量を検索するための API をサポートします。

Features are computed using batch computation or as streaming pipelines on Databricks and stored as Delta tables. These features can then be published using a scheduled Databricks Job or as a streaming pipeline. This ensures consistency of features used in batch training and features used in batch or online model inference, guaranteeing that there is no drift between features that were consumed at training and at serving time.

モデルのパッケージング

MLflow モデルフォーマットとの連携により、特徴量の情報は MLflow のモデルと一緒にパッケージングされます。Feature Registry と適切なプロバイダーから特徴量を検索できるように、Feature Store API はモデルのトレーニングの際、自動的にモデルのアーティファクト、全ての特徴量、実行時に用いられたコードをパッケージングします。

特徴量ストア情報を含む feature_spec.yaml は MLflow モデルアーティファクトと一緒にパッケージングされます。
特徴量ストア情報を含む feature_spec.yaml は MLflow モデルアーティファクトと一緒にパッケージングされます。

モデルアーティファクトと一緒にパッケージングされた特徴量情報によって、Feature Store のスコアリング API が入力データをスコアリングするために、格納されたキーを用いて自動的に特徴量を取得して結合します。このことは、モデルのデプロイメントの際にモデルを呼び出すクライアントは Feature Store とやり取りする必要がないことを意味します。結果として、クライアントを変更することなしに特徴量を更新できることになります。

Feature Store API ワークフロー

FeatureStoreClient Python ライブラリは、特徴量の計算・格納方法の定義、モデルトレーニングにおける既存特徴量の利用、バッチスコアリングにおける自動検索、オンラインストアへの特徴量の公開を行うために、Feature Store のコンポーネントとやりとりするための API を提供します。このライブラリは Databricks ML ランタイム(バージョン 8.3 以降)に同梱されています。

新規特徴量の作成

低レベル API は、カスタムの特徴量計算コードを記述するのに便利な機構を提供します。データサインティストは、ソーステーブルやファイルを用いて特徴量を計算するための Python の関数を記述します。

def compute_customer_features(data):
  '''Custom function to compute features and return a Spark DataFrame'''
  pass

customer_features_df = compute_customer_features(input_df)

新たな特徴量を作成して Feature Registry に登録するためには、create_feature_table API を呼び出します。特徴量を格納するための特定のデータベースやテーブルを指定することもできます。それぞれの特徴量テーブルは、エンティティの特徴量の値を識別するための主キーを持つ必要があります。

from databricks.feature_store import FeatureStoreClient

fs = FeatureStoreClient()

customer_feature_table = fs.create_feature_table(
  name='recommender.customer_features',
  keys='customer_id',
  features_df=customer_features_df,
  description='Customer features. Uses data from app interactions.'
)

Feature Store の特徴量を用いたモデルのトレーニング

Feature Store から特徴量を利用するためには、それぞれの特徴量テーブルが必要とする特徴量を識別するトレーニングセットを作成し、特徴量を検索し特徴量テーブルと結合するためのトレーニングデータセットのキーを記述します。

以下の例では、customer_features テーブルから二つの特徴量(total_purchases_30dpage_visits_7d)を使用し、トレーニングデータセットの customer_id を使用し、特徴量テーブルの主キーと突合して結合を行います。さらに、product_features テーブルの quantity_sold を使用し、product_id と country_code の複合主キーを用いて特徴量の検索、結合を行います。そして、create_training_set API を用いてトレーニングセットを構築し、モデルのトレーニングに不要なキーを削除します。

from databricks.feature_store import FeatureLookup

feature_lookups = [
    FeatureLookup(
      table_name = 'prod.customer_features',
      feature_name = 'total_purchases_30d',
      lookup_key = 'customer_id'
    ),
    FeatureLookup(
      table_name = 'prod.customer_features',
      feature_name = 'page_visits_7d',
      lookup_key = 'customer_id'
    ),
    FeatureLookup(
      table_name = 'prod.product_features',
      feature_name = 'quantity_sold',
      lookup_key = ['product_id', 'country_code']
    )
  ]

fs = FeatureStoreClient()

training_set = fs.create_training_set(
  df,
  feature_lookups = feature_lookups,
  label = 'rating',
  exclude_columns = ['customer_id', 'product_id']
)

モデルをトレーニングするためにあらゆる ML フレームワークが活用できます。log_model API は特徴量検索情報と ML モデルをパッケージングします。この情報はモデル推論の際に特徴量を検索するために用いられます。

from sklearn import linear_model

# Augment features specified in the training set specification
training_df = training_set.load_df().toPandas()

X_train = training_df.drop(['rating'], axis=1)
y_train = training_df.rating

model = linear_model.LinearRegression().fit(X_train, y_train)

fs.log_model(
  model,
  "recommendation_model",
  flavor=mlflow.sklearn,
  training_set=training_set,
  registered_model_name="recommendation_model"
)

特徴量の検索とパッチスコアリング

特徴量情報がパッケージングされたモデルをスコアリングする際、これらの特徴量はスコアリングの際に自動的に検索されます。ユーザーが指定する必要があるのは、入力データの列を検索するための主キーです。Feature Store の score_batch API は内部でモデルアーティファクトに格納された特徴量の仕様を用いて、特定のテーブル、特徴量の列、結合キーに関して Feature Registry とやりとりを行います。そして、この API は、モデルのスコアリングに必要なスキーマのデータフレームを作成するために、適切な特徴量テーブルと効率的な結合を行います。以下のシンプルなコードでは、上でトレーニングしたモデルに関するオペレーションを説明しています。

# batch_df has columns ‘customer_id’ and ‘product_id’
predictions = fs.score_batch(
    model_uri,
    batch_df
)

# The returned ‘predictions’ dataframe has these columns:
#  inputs from batch_df: ‘customer_id’, ‘product_id’
#  features: ‘total_purchases_30d’, ‘page_visits_7d’, ‘quantity_sold’
#  model output: ‘prediction’

オンラインストアでの特徴量の公開

オンラインストアに特徴量テーブルを公開するには、まず初めにオンランストアのスペックを指定し、 publish_table API を呼び出します。以降のコードは、バッチプロバイダーの customer_features テーブルを最新の特徴量で既存のテーブルを上書きします。

online_store = AmazonRdsMySqlSpec(hostname, port, user, password)

fs.publish_table(
  name='recommender_system.customer_features',
  online_store=online_store,
  mode='overwrite'
)

publish_features は特定の特徴量の値(日付などのフィルタリング条件)を除外するためのさまざまなオプションをサポートしています。以下の例では、今日時点の特徴量をオンラインストアに流し込み、既存の特徴量とマージします。

fs.publish_table(
  name='recommender_system.customer_features',
  online_store=online_store,
  filter_condition=f"_dt = '{str(datetime.date.today())}'",
  mode='merge',
  streaming='true'
)

Feature Store(特徴量ストア)を使い始める

Ready to get started or try it out for yourself? You can read more about Databricks Feature Store and how to use it in our documentation at AWS, Azure and GCP.

Feature Store (特徴量ストア)についての詳細は、「データ+ AI サミット」の基調講演の動画で無料でご覧になれます。
Databricks 無料トライアル 使ってみる

ご登録

Feature Store (特徴量ストア)についての詳細は、「データ+ AI サミット」の基調講演の動画で無料でご覧になれます。