Centralizing ML Features through Feature Store in Google Cloud Vertex AI
Imagine you are the head chef. Your primary responsibility is to cook the food, and prepare it in a way that is appetizing to the customers. The moment you were about to start cooking, you notice that the ingredients were still not prepared. The carrots and onions were still not peeled, the meat is still frozen, and the other ingredients in general were not ready. If only there was a sous chef that would help prepare all the ingredients so that the head chef would only need to cook the food.
In a way, this specific part of the cooking process is analogous to the pre-training segment of the Machine Learning (ML) workflow. The raw ingredients would correspond to the raw and untouched data. Same as cooking, the ML Engineer cannot use the raw unprepared data to train a model and expect good results. There needs to be preparation beforehand, which concerns data preprocessing then feature engineering. You can learn more about both of them in my article Data Preprocessing with Google Cloud Dataprep by Trifacta. The product of featuring engineering would be these so called features.
What is a Feature?
In the context of ML, a feature refers to any individual measurable property or characteristic of the data that is fed into an ML model for both training and inference. Features are the input variables that the model uses to learn patterns during the training phase and to make predictions during the inference phase.
During training, the model is exposed to historical data, where each feature helps the model to identify and learn underlying patterns or relationships within the data. These patterns are then used by the model to make predictions on new, unseen data during the serving phase. The quality and relevance of the features directly influence the model’s performance, accuracy, and generalization to new data.
Feature Engineering
Given the critical role of features in model performance, careful selection and transformation of features, a process known as feature engineering, is essential. Feature engineering involves creating, modifying, and selecting features that improve the predictive power of the model. This process can include techniques such as scaling, encoding categorical variables, creating interaction terms, and extracting useful information from raw data. The goal of feature engineering is to enhance the model’s ability to capture the true signal in the data while minimizing noise, ultimately leading to better predictions and more robust models.
Feature Store
Now that we’ve prepared the ingredients properly through feature engineering. We need a place to store the features that we’ve preprocessed. The place where we store features is usually called a feature store in ML System Design. It serves as a hub where all features are gathered together in order to have a single source of truth across multiple platforms. By putting the features in a single place, we can make sure that everyone is using the same quality of feature data, forged by the same transformations.
In many organizations, different teams often work in isolation, each developing models that rely on similar features. However, these features are frequently defined in varied ways and computed independently, resulting in redundant efforts and increased computing and storage expenses. A feature store serves as a centralized hub that standardizes ML features across the organization. With a feature store, one ML team can seamlessly utilize features created by another team. Since these features are cataloged, shared, and pre-computed, the organization can significantly reduce cloud costs by eliminating the need to recompute the same features repeatedly.
Components of a Feature Store
Feature Transformation
This is generally under feature engineering or preprocessing but some features store services include this as the initial step, where raw data is preprocessed into usable features. One of the core functionalities of a feature store is its ability to store not only features, but the logic as to how raw data was transformed into usable data. This ensure that all raw data coming from different sources are preprocessed using the same transformations, preventing training-serving skew. Training-serving skew refers to a difference between training and serving features. This issue is usually caused by a discrepancy in managing transformations in training and serving pipelines, a change in data between training and production, etc. Through proper setup of the transformation pipeline during feature engineering, we can avoid this detrimental issue.
Feature Storage
This had been already mentioned, but to further expound, the feature storage component of a features store is where the processed features are stored after the transformations. Feature storage can be thought of as a dual-database system, which corresponds to the two types of features it usually holds: Offline and Online stores.
First is the static Offline Store, which houses historical feature data used for training ML models. It typically stores large volumes of data in a way that allows for batch processing and querying, enabling data scientists to access and use historical feature values when building and iterating on models. The offline store is optimized for high-throughput and large-scale data retrieval, often integrating with data lakes or distributed file systems, making it a crucial component for model development and evaluation.
An Online Store on the other hand is a low-latency, high-availability storage system designed to serve real-time feature data for model inference. It stores the most recent and relevant feature values, enabling ML models to make fast predictions based on current data. The online store is optimized for quick look-ups and updates, ensuring that the features served during model inference are fresh and consistent with the real-time state of the system. This component is critical for applications that require immediate predictions, such as recommendation engines or fraud detection systems.
Feature Registry
The feature registry is a crucial component of a feature store that acts as a centralized catalog for managing, organizing, and discovering features. It stores metadata about each feature, including its definition, source, data type, and lineage, providing a unified view of all available features across teams and projects. The registry ensures consistency and re-usability of features by enabling data scientists and engineers to easily find and share feature definitions, reducing redundancy and improving collaboration. By maintaining a comprehensive record of features, the registry also facilitates governance, versioning, and monitoring, ensuring that features used in production are well-documented and compliant with organizational standards.
Feature Serving
Feature serving is a critical component of a feature store that delivers features to ML models during inference, and it comes in two primary modes: Batch and Online serving. Batch serving involves delivering features to models in large, precomputed sets, typically used for offline processes like model training or periodic batch predictions. In contrast, Online serving focuses on real-time feature delivery, where features are retrieved from the online store, potentially transformed on the fly, and fed into the model with minimal latency. While batch serving is optimized for throughput and scale, online serving is tailored for speed and low latency, ensuring that time-sensitive applications can make quick and accurate predictions based on the latest data.
Vertex AI Feature Store
Google Cloud has its own take on feature stores, and they simply call it Feature Store. It is integrated within the Vertex AI ecosystem, therefore it can be easily connected with other phases of the ML workflow pipeline. However, feature engineering is not inherently available to Vertex AI Feature Store. You need to use external data preprocessing services, but in this case, you can just use Dataprep by Trifacta. Dataprep takes in raw data, which can be taken from BigQuery, and it outputs the transformed features also back in BigQuery. Conveniently, we can take feature data from BigQuery and feed it into Feature Store.
Getting Started
Vertex AI Feature Store can be accessed under the Vertex AI tab under the Data category within the Google Cloud Platform console. Before going to Feature Store, you need to already have a BigQuery table, for it is where the feature data is extracted from. For this article I used the Daily Temperature of Major Cities from Kaggle, where daily average temperature data from 1995 were recorded.
To centralize your features, you first need to create a feature group. A feature group is a logical grouping of features which are related to the same entity or concept. It is essentially a collection of features that are managed together, typically because they are associated with a common entity. To create a feature group, go the the feature group tab under Feature Store, and select create.
Fill out the information needed, including the BigQuery path where your feature data is initially stored, and the columns (which represent the features).
You can then go back to the Feature Store dashboard to see the feature group that you created.
The features are now stored somewhere in the cloud, and had been centralized in a hub. The next step is to send these features somewhere for training or inference, which is called feature serving. There are two types of feature serving, one for transferring bulk training data, and the other for feeding features for inference, which require low latency. Online serving in general refers to the process of providing low-latency access to ML features for real-time predictions or inferences. This capability is essential for use cases where the ML model needs to make predictions in real time. Vertex AI Feature Store offers two serving types, Bigtable online serving and Optimized online serving. As you may have noticed, Vertex AI calls both general types of serving variants of “Online serving”, which is quite misleading. Bigtable online serving can be used for transferring large volumes of data, which means higher latency, but is more cost-efficient in the greater scheme. This is suitable for serving features that will be used as training data. Optimized online serving on the other hand serves features at a much lower latency. This means that it is suitable in scenarios where it’s critical to serve features for inference at ultra-low latency.
To serve features from a feature group using Vertex AI, you need to create an online store. This can be accessed under the online store tab in Feature Store, and then by selecting create.
You need to fill out the details needed, where you get to choose the service, either Bigtable or Optimized online store.
After creating the online store, you need some sort of interface to serve your features, and this way you can merge other features from different online stores. This interface is called a feature view. A feature view is also a logical entity that defines a specific way of organizing and accessing a set of features from the feature registry. It is essentially a structured schema that groups together related features, specifying how these features should be retrieved and used in both training and serving models. The concept is quite similar to a feature group, but make sure to differentiate the two, as feature view is a schema that translates into an interface, allowing different combinations of features from the registry.
To create a feature view under an online store for serving, select the three vertical dots under a single instance of the online store, and then select create feature view.
When creating a feature view, you can opt to modify the sync config, which is an optional configuration that you can specify when creating a feature view to manage how the data is synchronized between the offline store and the online store. The offline store in this case refers to the feature group from the Feature Registry or a table from BigQuery.
With this, you have essentially exposed the feature view as an API, which can be accessed through an external call. You can read more about the this in the Vertex AI Feature Store documentation, but in summary, this is how you would access the feature view in Python using the Google Cloud Software Development Kit (directly from the documentation).
from google.cloud.aiplatform_v1 import FeatureOnlineStoreServiceClient
from google.cloud.aiplatform_v1.types import feature_online_store_service as feature_online_store_service_pb2
data_client = FeatureOnlineStoreServiceClient(
client_options={
"api_endpoint": f"LOCATION_ID-aiplatform.googleapis.com"
}
)
data_client.fetch_feature_values(
request=feature_online_store_service_pb2.FetchFeatureValuesRequest(
feature_view=f"projects/PROJECT_ID/locations/LOCATION_ID/featureOnlineStores/FEATUREONLINESTORE_NAME/featureViews/FEATUREVIEW_NAME",
data_key=feature_online_store_service_pb2.FeatureViewDataKey(key="ENTITY_ID"),
data_format=feature_online_store_service_pb2.FeatureViewDataFormat.FORMAT
)
)Replace the following:
- LOCATION_ID: Region where the online store is located, such as `us-central1`.
- PROJECT_ID: Your project ID.
- FEATUREONLINESTORE_NAME: The name of the online store containing the feature view.
- FEATUREVIEW_NAME: The name of the feature view from which you want to serve feature values.
- ENTITY_ID: The value of the ID column in the feature record from which you want to serve the latest feature values.
- FORMAT: Optional: The format in which you want to fetch the feature values. Supported formats include JSON `KEY_VALUE` pair and proto `PROTO_STRUCT` formats.
What’s Next?
In this article, we have stored preprocessed features from BigQuery into Vertex AI Feature Store as a feature group. We also served the stored features as a feature view within online stores, and accessed the endpoint through an API call using the Google Cloud SDK in Python. This means that you can integrate the feature data programmatically, and feed it into you model. Through this, we have completed the pre-training stages of the ML workflow, from raw data to usable features. The next stage in the workflow after training is developing or building the ML model and fitting the model to the data. It is in this phase where the features are fed into the model so that it could learn the data patterns. Google Cloud provides services for building or developing a model, and this is highlighted in my article AI Development and Integration Using Google Cloud, which I recommend as your next reading material.
References
Duta, G. (2024, June). What is a Feature Store in ML, and Do I Need One?. Qwak. https://www.qwak.com/post/what-is-a-feature-store-in-ml
Google Cloud Training. (n.d.). Feature Engineering [MOOC]. Coursera. https://www.coursera.org/learn/feature-engineering
Google Cloud. (n.d.). Serve features from online store. https://cloud.google.com/vertex-ai/docs/featurestore/latest/serve-feature-values
