Vertex AI Feature Storeのオフラインストア機能を活用した特徴量管理術

はじめに

こんにちは！AI技術開発室のYTです。
機械学習プロジェクトにおいて、特徴量の管理はモデルの性能や開発効率を左右する重要な要素です。しかし、「あのモデルで使った特徴量はどこだっけ？」「学習時と推論時で特徴量の意味合いが変わってしまった」といった課題に直面することも少なくありません。本記事では、こうした課題を解決する強力なツールとして、Google CloudのVertex AI Feature Storeのオフラインストア機能をご紹介します。

Vertex AI Feature Store入門

Vertex AI Feature Storeは、Google Cloudが提供する、機械学習（ML）用の特徴量を一元的に整理・保存し、提供するフルマネージドサービスです。
特徴量の定義やメタデータを管理する「特徴レジストリ(Feature Registry)」の機能を含んでおり、特徴量のカタログ化、バージョン管理、そして学習や推論のための効率的なサービング（提供）を可能にし、MLワークフロー全体での特徴量の一貫性と再利用性を高めます。これにより、開発サイクルの短縮、モデル品質の向上、MLOpsにおけるガバナンス強化に貢献します。

オフラインストアとオンラインストアの違い

Vertex AI Feature Storeには、モデルへの特徴量提供の方法として「オンラインストア」と「オフラインストア」という2種類のサービング形式があり、それぞれ以下のような特性があります。なお、オフラインストアについては、BigQueryのテーブル/ビューがその役割を果たすため、Vertex AI内に別途ストアを構築する必要はありません。

特性	オフラインストア	オンラインストア
用途	モデル学習用の大規模な履歴特徴量の保存・管理、特徴量の探索・分析、バッチ推論、Point-in-Time Correctnessの実現、特徴モニタリングによる特徴量ドリフトの検出	リアルタイム推論のための低レイテンシな最新特徴量提供
データ	過去数年分など、大規模な時系列化された特徴量	最新の、あるいは頻繁に更新される特徴量のスナップショット
レイテンシ	秒〜分単位 (バッチ処理向け)	ミリ秒単位
主な技術	BigQuery	Cloud Bigtableなど
更新頻度	バッチ処理で定期的 (例: 日次、時間次)	ほぼリアルタイム更新可能
コスト	BigQuery利用料ベースで比較的安価	サービングストレージ量やノード数等に依存し、比較的高価

Vertex AIには、類似する機能としてマネージドデータセットもありますが、こちらは主にAutoML時に利用します。本記事ではBigQueryベースのオフラインストア機能についてご紹介させていただき、レガシー版のオフラインストアやオンラインストア機能についてはスコープ外とさせていただきます。

なぜBigQueryそのものではなくオフラインストアに登録するのか？

「オフラインストアの実体はBigQueryなら、直接BigQueryを使えば良いのでは？」と感じるかもしれません。しかし、Feature Storeのオフラインストアを利用することには、単にBigQueryを直接操作する以上のメリットがあります。

• 特徴量管理の一元化と標準化: Feature Storeを中心的な特徴量カタログとして利用することで、チーム内や組織全体でどのような特徴量が利用可能かを発見しやすくなり、特徴量の再利用を促進します。
• Point-in-Time Correctnessの容易な実現: Vertex AI Feature Storeを活用することで、各エンティティの特定時点での特徴量を正確に取得できます。これにより、信頼性の高い学習データを効率的に準備できます。
• オンラインストアへのシームレスな連携: オフラインストアで検証・準備した特徴量の一部または全部を、簡単な設定でオンラインストアに同期し、リアルタイム推論に活用できます。この連携のスムーズさは大きな利点です。
• 特徴データの異常検出の容易さ: 記事執筆時点ではプレビューですが、時間の経過とともに特徴データに発生するドリフトを検出する特徴モニタリング機能が提供されています。
• Vertex AI エコシステムとの統合: Vertex AI Feature Storeは、MLOpsパイプラインの効率化に貢献します。特に Vertex AI Pipelinesと連携し、特徴量関連のワークフロー自動化を支援します。また、Feature Storeで一元管理された特徴量定義とメタデータは、Vertex AI Trainingで用いる学習データの一貫性確保や、Vertex AI Model Registryに登録されたモデルのトレーサビリティ向上に繋がり、MLOps全体のガバナンス強化を後押しします。

オフラインストア機能を構成する主なコンセプト

• 特徴グループ(Feature Group): 個々の特徴量をグループ化したもので、BigQueryの1つのテーブル/ビューに相当します。後述するエンティティID、特徴、特徴タイムスタンプが含まれます。
• 特徴(Feature): モデルで使用される個々の入力変数であり、BigQueryテーブル/ビュー内の各カラムに相当します。
• エンティティID: 特徴量が紐づく対象（例：ユーザーや商品）を一意に識別するキーで、主キーに相当するカラムです。特徴取得や結合の際に基準となります。
• 特徴タイムスタンプ: 各特徴量の値が生成または観測された時刻を示すタイムスタンプです。Point-in-Time Correctnessを実現するための非常に重要な要素で、過去の特定の時点での特徴量の状態を正確に再現でき、学習データ生成時に未来の情報がリークするのを防ぎます。

購買データを例にすると、以下のようになります。

オフラインストアへの特徴量登録

1. ソースとなるBigQueryテーブル/ビューの設計と考慮事項

オフラインストアに特徴量を登録するには、まず BigQuery上に特徴量データを含むテーブルまたはビューを準備し、それをVertex AI Feature Storeのリソース(特徴グループ, 特徴)に紐づけます。ソースとなるBigQueryテーブル/ビューには、上記で解説したエンティティIDカラムや特徴カラム、特徴タイムスタンプカラムを含める必要があります。詳しくはこちらのガイドラインをご参照ください。

BigQueryテーブルの最適化

BigQueryのTIMESTAMP型はデフォルトでUTCを基準とするため、特徴タイムスタンプカラムもUTCで統一して管理することで、タイムゾーンの不一致によるクエリ結果のズレを防止できます。
また、特徴タイムスタンプカラムでパーティショニングを行うことにより、特定の期間のデータ取得や学習データセット生成時のクエリスキャン量を大幅に削減し、パフォーマンス向上とコスト削減に繋がります。エンティティIDカラムや頻繁にフィルタ条件として使用する特徴量カラムでクラスタリングを行うと、クエリパフォーマンスが向上する場合があります。

2. 特徴グループの作成

特徴グループを作成するには以下のPythonコードを実行します。コンソール画面からの作成も可能です。

from google.cloud import aiplatform
from google.cloud.aiplatform_v1 import FeatureRegistryServiceClient
from google.cloud.aiplatform_v1.types import feature_group as feature_group_pb2
from google.cloud.aiplatform_v1.types import feature_registry_service as feature_registry_service_pb2
from google.cloud.aiplatform_v1.types import io as io_pb2

PROJECT_ID = "your_project" # Google Cloudプロジェクト名
LOCATION = "your_location" # Region asia-northeast1など
FEATURE_GROUP_ID = "your_feature_group_id" # 作成する特徴グループ名
DATASET_ID = "your_dataset" # ソースのBQテーブルのデータセット名
TABLE_NAME = "your_feature_group_table" # ソースのBQテーブルorビュー名
ENTITY_ID_COL = "entity_id" # エンティティID列として利用する列名

aiplatform.init(project=PROJECT_ID, location=LOCATION)

registry_client = FeatureRegistryServiceClient(
    client_options={"api_endpoint": f"{LOCATION}-aiplatform.googleapis.com"}
)

feature_group_config = feature_group_pb2.FeatureGroup(
    big_query=feature_group_pb2.FeatureGroup.BigQuery(
        big_query_source=io_pb2.BigQuerySource(input_uri=f"{PROJECT_ID}.{DATASET_ID}.{TABLE_NAME}"),
        entity_id_columns=[ENTITY_ID_COL],
    ),
    description="This is a feature group for testing",
)

create_group_lro = registry_client.create_feature_group(
    feature_registry_service_pb2.CreateFeatureGroupRequest(
        parent=f"projects/{PROJECT_ID}/locations/{LOCATION}",
        feature_group_id=FEATURE_GROUP_ID,
        feature_group=feature_group_config,
    )
)
create_group_lro.result()

3. 特徴の登録

特徴グループの作成に引き続き、特徴を作成するには、以下のPythonコードを実行します。コンソール画面からの作成も可能です。

from google.cloud.aiplatform_v1.types import featurestore_service as featurestore_service_pb2
FEATURE_IDS = ["feature_name_1", "feature_name_2"] # 特徴として登録する列名リスト

create_feature_lros = []
for fid in FEATURE_IDS:
    create_feature_lros.append(
        registry_client.create_feature(
            featurestore_service_pb2.CreateFeatureRequest(
          parent=f"projects/{PROJECT_ID}/locations/{LOCATION}/featureGroups/{FEATURE_GROUP_ID}",
                feature_id=fid,
                feature=feature_pb2.Feature(),
            )
        )
    )
for lro in create_feature_lros:
    lro.result()

オフラインストアからの特徴量取得

オフラインストアから特徴量を取得する際にはBigQueryのクライアントをそのまま利用できます。以下によくあるデータ取得パターンに対応するクエリ例を示します。

最新のデータを取得する場合 (各エンティティの最新値)

SELECT
  * -- 必要な列を選択
FROM
  `your_project.your_dataset.your_feature_group_table`
QUALIFY
  ROW_NUMBER() OVER (PARTITION BY entity_id ORDER BY feature_timestamp DESC) = 1;

2025年4月1日午前0時時点のデータを取得する場合 (Point-in-time)

SELECT
  * -- 必要な列を選択
FROM
  `your_project.your_dataset.your_feature_group_table`
WHERE
  -- 指定時刻以前のデータに絞り込む
  feature_timestamp <= TIMESTAMP('2025-04-01 00:00:00')
QUALIFY
  -- 各エンティティで、指定時刻以前の最新のものを選択
  ROW_NUMBER() OVER (PARTITION BY entity_id ORDER BY feature_timestamp DESC) = 1;

まとめ

今回は、Vertex AI Feature Storeのオフラインストア機能と、その活用方法の一例をご紹介させていただきました。現時点ではオンラインストアの利用を計画されていなくても、オフラインストア機能を活用することで、特徴量管理の作業が効率化されたり、より質の高い特徴量エンジニアリングへの道が開けたりするかもしれません。

We are hiring!

AI技術開発室では、一緒に働く仲間を募集しています。技術を追究しながら「知と情熱と行動力で、医療人を支え、医療の未来を動かす。」というパーパスをともに実現しましょう。詳しくは以下をご覧ください。
https://hrmos.co/pages/carenet5800/jobs/0000020