Matchmaking with Federated Learning: The Future of Privacy-Centric Dating Apps (Startup code) - Part III

In the "near" real-world implementation of "Approach 2 - Incorporating User Features into a Single Model", for initial training and Federated Learning using TensorFlow and TensorFlow Federated (TFF). This example will focus on a recommendation system for our dating app, where the model predicts user preferences for faces based on extracted features and user embeddings.

Initial Setup

Install Dependencies:

pip install tensorflow tensorflow-federated        

Import Libraries:

import tensorflow as tf
import tensorflow_federated as tff
import numpy as np
from tensorflow import keras
from tensorflow.keras import layers        

tensorflow_federated: The TensorFlow Federated library for simulating and implementing federated learning.

Data Preparation

Assuming you have a dataset of user interactions with face images, where each interaction includes user ID, extracted face features, and a binary label indicating like/dislike.

Load and pre-process your data:

This block simulates a dataset for demonstration purposes. In a real-world scenario, you would load your actual dataset.

# Load your dataset
# For this example, we'll create a dummy dataset
num_users = 100
num_faces = 1000
face_feature_dim = 128

# Dummy data: user IDs, face features, and labels (like/dislike)
user_ids = np.random.randint(0, num_users, size=(num_faces,))
face_features = np.random.rand(num_faces, face_feature_dim)
labels = np.random.randint(0, 2, size=(num_faces,))

# Split data into training and testing
split_index = int(num_faces * 0.8)
train_data = (user_ids[:split_index], face_features[:split_index], labels[:split_index])
test_data = (user_ids[split_index:], face_features[split_index:], labels[split_index:])
        

'num_users', 'num_faces', and 'face_feature_dim' are constants representing the number of users, number of face images, and dimensionality of face features, respectively.

'user_ids', 'face_features', and 'labels' are randomly generated arrays representing user IDs, face features, and like/dislike labels for each face.

Model Preparation

Define Recommendation model:

class RecommendationModel(keras.Model):
    def __init__(self, num_users, embedding_dim, face_feature_dim):
        super(RecommendationModel, self).__init__()
        self.user_embedding = layers.Embedding(input_dim=num_users, output_dim=embedding_dim)
        self.fc_layers = keras.Sequential([
            layers.Dense(128, activation='relu'),
            layers.Dropout(0.5),
            layers.Dense(64, activation='relu'),
            layers.Dropout(0.5),
            layers.Dense(1, activation='sigmoid')
        ])

    def call(self, inputs):
        user_ids, face_features = inputs
        user_embedding = self.user_embedding(user_ids)
        combined_features = tf.concat([user_embedding, face_features], axis=1)
        return self.fc_layers(combined_features)
        

• RecommendationModel is a subclass of keras.Model and represents the neural network model for predicting user preferences.
• user_embedding is an embedding layer that maps user IDs to embedding vectors representing user preferences.
• fc_layers is a sequential model consisting of fully connected (Dense) layers with ReLU activation and Dropout for regularization. The final layer uses a Sigmoid activation function to output a probability between 0 and 1, representing the likelihood of a user liking a face.
• The call method defines the forward pass of the model. It takes user_ids and face_features as inputs, combines the user embedding with the face features, and passes the combined features through the fully connected layers to produce the output.

Initial Training

Train the model:

embedding_dim = 50

model = RecommendationModel(num_users, embedding_dim, face_feature_dim)
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# Convert the data to a TensorFlow dataset and batch it
train_dataset = tf.data.Dataset.from_tensor_slices(train_data).batch(32)
test_dataset = tf.data.Dataset.from_tensor_slices(test_data).batch(32)

# Train the model
model.fit(train_dataset, epochs=10, validation_data=test_dataset)        

• embedding_dim is the dimensionality of the user embedding vectors.
• model is an instance of the RecommendationModel class.
• model.compile compiles the model with the Adam optimizer, binary crossentropy loss function (suitable for binary classification), and accuracy as a metric.

Federated Learning Setup

For Federated Learning, we'll use TensorFlow Federated (TFF) to simulate a federated environment.

Define Federated Data and Model:

# Define a TFF computation for creating a federated dataset
def create_federated_dataset(client_data):
    return [tf.data.Dataset.from_tensor_slices(data).batch(32) for data in client_data]

# Split the training data into federated data for two clients
federated_train_data = create_federated_dataset([train_data[:split_index // 2], train_data[split_index // 2:]])

# Wrap the Keras model with TFF
def model_fn():
    return tff.learning.from_keras_model(
        model,
        input_spec=train_dataset.element_spec,
        loss=tf.keras.losses.BinaryCrossentropy(),
        metrics=[tf.keras.metrics.Accuracy()]
    )        

• 'create_federated_dataset' is a function that takes client data and creates a federated dataset by splitting the data into batches for each client.
• 'federated_train_data' is the federated dataset created from the training data, split between two clients for simulation.
• 'model_fn' is a function that returns a TFF-wrapped Keras model. It specifies the model, input specification (derived from the training dataset), loss function, and metrics for federated learning.

Federated Training:

# Define the federated averaging process
iterative_process = tff.learning.build_federated_averaging_process(model_fn)

# Initialize the federated averaging process
state = iterative_process.initialize()

# Run the federated averaging process for a number of rounds
num_rounds = 10
for round_num in range(1, num_rounds + 1):
    state, metrics = iterative_process.next(state, federated_train_data)
    print(f'Round {round_num}, Metrics: {metrics}')        

• iterative_process is a TFF iterative process that represents the federated averaging algorithm. It is built using the model_fn function.
• state is the initial state of the federated averaging process.
• The for loop simulates federated learning rounds. In each round, the model is trained on the federated dataset, and the aggregated updates are applied to the global model. Metrics are printed after each round.

In this implementation:

• The RecommendationModel class defines a neural network that combines user embeddings with face features to predict user preferences.
• For initial training, we use the standard Keras API to train the model on a centralized dataset.
• For Federated Learning, we use TensorFlow Federated (TFF) to simulate a federated environment. We define a model_fn function that wraps our Keras model for use with TFF, and we use the tff.learning.build_federated_averaging_process to create a federated averaging process. This process is then executed for a number of rounds, with the model being trained collaboratively by multiple clients (simulated by splitting the training data).

Here are some of the learning resources that have helped me master the topic -

• TensorFlow Federated Tutorials by Google
• TensorFlow Federated Tutorial Session by Google TechTalks
• Introduction to Remote Data Science by Openmined.org
• Secure and Private AI by Udacity in collaboration with authors at Openmined

In next article we would talk about Secure Aggregation, Robustness and Fault tolerance, Scalability, Model and Data Versioning, Monitoring and Logging, Integration with Mobile devices, and Compliance and Ethics.