Memory-efficient embeddings for recommendation systems

import os os.environ["KERAS_BACKEND"] = "tensorflow" from zipfile import ZipFile from urllib.request import urlretrieve import numpy as np import pandas as pd import tensorflow as tf import keras from keras import layers from keras.layers import StringLookup import matplotlib.pyplot as plt 

urlretrieve("http://files.grouplens.org/datasets/movielens/ml-1m.zip", "movielens.zip") ZipFile("movielens.zip", "r").extractall() ratings_data = pd.read_csv( "ml-1m/ratings.dat", sep="::", names=["user_id", "movie_id", "rating", "unix_timestamp"], ) ratings_data["movie_id"] = ratings_data["movie_id"].apply(lambda x: f"movie_{x}") ratings_data["user_id"] = ratings_data["user_id"].apply(lambda x: f"user_{x}") ratings_data["rating"] = ratings_data["rating"].apply(lambda x: float(x)) del ratings_data["unix_timestamp"] print(f"Number of users: {len(ratings_data.user_id.unique())}") print(f"Number of movies: {len(ratings_data.movie_id.unique())}") print(f"Number of ratings: {len(ratings_data.index)}") 

random_selection = np.random.rand(len(ratings_data.index)) <= 0.85 train_data = ratings_data[random_selection] eval_data = ratings_data[~random_selection] train_data.to_csv("train_data.csv", index=False, sep="|", header=False) eval_data.to_csv("eval_data.csv", index=False, sep="|", header=False) print(f"Train data split: {len(train_data.index)}") print(f"Eval data split: {len(eval_data.index)}") print("Train and eval data files are saved.") 

csv_header = list(ratings_data.columns) user_vocabulary = list(ratings_data.user_id.unique()) movie_vocabulary = list(ratings_data.movie_id.unique()) target_feature_name = "rating" learning_rate = 0.001 batch_size = 128 num_epochs = 3 base_embedding_dim = 64 

def get_dataset_from_csv(csv_file_path, batch_size=128, shuffle=True): return tf.data.experimental.make_csv_dataset( csv_file_path, batch_size=batch_size, column_names=csv_header, label_name=target_feature_name, num_epochs=1, header=False, field_delim="|", shuffle=shuffle, ) def run_experiment(model): # Compile the model. model.compile( optimizer=keras.optimizers.Adam(learning_rate), loss=keras.losses.MeanSquaredError(), metrics=[keras.metrics.MeanAbsoluteError(name="mae")], ) # Read the training data. train_dataset = get_dataset_from_csv("train_data.csv", batch_size) # Read the test data. eval_dataset = get_dataset_from_csv("eval_data.csv", batch_size, shuffle=False) # Fit the model with the training data. history = model.fit( train_dataset, epochs=num_epochs, validation_data=eval_dataset, ) return history 

def embedding_encoder(vocabulary, embedding_dim, num_oov_indices=0, name=None): return keras.Sequential( [ StringLookup( vocabulary=vocabulary, mask_token=None, num_oov_indices=num_oov_indices ), layers.Embedding( input_dim=len(vocabulary) + num_oov_indices, output_dim=embedding_dim ), ], name=f"{name}_embedding" if name else None, ) 

def create_baseline_model(): # Receive the user as an input. user_input = layers.Input(name="user_id", shape=(), dtype=tf.string) # Get user embedding. user_embedding = embedding_encoder( vocabulary=user_vocabulary, embedding_dim=base_embedding_dim, name="user" )(user_input) # Receive the movie as an input. movie_input = layers.Input(name="movie_id", shape=(), dtype=tf.string) # Get embedding. movie_embedding = embedding_encoder( vocabulary=movie_vocabulary, embedding_dim=base_embedding_dim, name="movie" )(movie_input) # Compute dot product similarity between user and movie embeddings. logits = layers.Dot(axes=1, name="dot_similarity")( [user_embedding, movie_embedding] ) # Convert to rating scale. prediction = keras.activations.sigmoid(logits) * 5 # Create the model. model = keras.Model( inputs=[user_input, movie_input], outputs=prediction, name="baseline_model" ) return model baseline_model = create_baseline_model() baseline_model.summary() 

history = run_experiment(baseline_model) plt.plot(history.history["loss"]) plt.plot(history.history["val_loss"]) plt.title("model loss") plt.ylabel("loss") plt.xlabel("epoch") plt.legend(["train", "eval"], loc="upper left") plt.show() 

class QREmbedding(keras.layers.Layer): def __init__(self, vocabulary, embedding_dim, num_buckets, name=None): super().__init__(name=name) self.num_buckets = num_buckets self.index_lookup = StringLookup( vocabulary=vocabulary, mask_token=None, num_oov_indices=0 ) self.q_embeddings = layers.Embedding( num_buckets, embedding_dim, ) self.r_embeddings = layers.Embedding( num_buckets, embedding_dim, ) def call(self, inputs): # Get the item index. embedding_index = self.index_lookup(inputs) # Get the quotient index. quotient_index = tf.math.floordiv(embedding_index, self.num_buckets) # Get the reminder index. remainder_index = tf.math.floormod(embedding_index, self.num_buckets) # Lookup the quotient_embedding using the quotient_index. quotient_embedding = self.q_embeddings(quotient_index) # Lookup the remainder_embedding using the remainder_index. remainder_embedding = self.r_embeddings(remainder_index) # Use multiplication as a combiner operation return quotient_embedding * remainder_embedding 

class MDEmbedding(keras.layers.Layer): def __init__( self, blocks_vocabulary, blocks_embedding_dims, base_embedding_dim, name=None ): super().__init__(name=name) self.num_blocks = len(blocks_vocabulary) # Create vocab to block lookup. keys = [] values = [] for block_idx, block_vocab in enumerate(blocks_vocabulary): keys.extend(block_vocab) values.extend([block_idx] * len(block_vocab)) self.vocab_to_block = tf.lookup.StaticHashTable( tf.lookup.KeyValueTensorInitializer(keys, values), default_value=-1 ) self.block_embedding_encoders = [] self.block_embedding_projectors = [] # Create block embedding encoders and projectors. for idx in range(self.num_blocks): vocabulary = blocks_vocabulary[idx] embedding_dim = blocks_embedding_dims[idx] block_embedding_encoder = embedding_encoder( vocabulary, embedding_dim, num_oov_indices=1 ) self.block_embedding_encoders.append(block_embedding_encoder) if embedding_dim == base_embedding_dim: self.block_embedding_projectors.append(layers.Lambda(lambda x: x)) else: self.block_embedding_projectors.append( layers.Dense(units=base_embedding_dim) ) def call(self, inputs): # Get block index for each input item. block_indicies = self.vocab_to_block.lookup(inputs) # Initialize output embeddings to zeros. embeddings = tf.zeros(shape=(tf.shape(inputs)[0], base_embedding_dim)) # Generate embeddings from blocks. for idx in range(self.num_blocks): # Lookup embeddings from the current block. block_embeddings = self.block_embedding_encoders[idx](inputs) # Project embeddings to base_embedding_dim. block_embeddings = self.block_embedding_projectors[idx](block_embeddings) # Create a mask to filter out embeddings of items that do not belong to the current block. mask = tf.expand_dims(tf.cast(block_indicies == idx, tf.dtypes.float32), 1) # Set the embeddings for the items not belonging to the current block to zeros. block_embeddings = block_embeddings * mask # Add the block embeddings to the final embeddings. embeddings += block_embeddings return embeddings 

movie_frequencies = ratings_data["movie_id"].value_counts() movie_frequencies.hist(bins=10) 

sorted_movie_vocabulary = list(movie_frequencies.keys()) movie_blocks_vocabulary = [ sorted_movie_vocabulary[:400], # high popularity movies block sorted_movie_vocabulary[400:1700], # normal popularity movies block sorted_movie_vocabulary[1700:], # low popularity movies block ] movie_blocks_embedding_dims = [64, 32, 16] user_embedding_num_buckets = len(user_vocabulary) // 50 def create_memory_efficient_model(): # Take the user as an input. user_input = layers.Input(name="user_id", shape=(), dtype="string") # Get user embedding. user_embedding = QREmbedding( vocabulary=user_vocabulary, embedding_dim=base_embedding_dim, num_buckets=user_embedding_num_buckets, name="user_embedding", )(user_input) # Take the movie as an input. movie_input = layers.Input(name="movie_id", shape=(), dtype="string") # Get embedding. movie_embedding = MDEmbedding( blocks_vocabulary=movie_blocks_vocabulary, blocks_embedding_dims=movie_blocks_embedding_dims, base_embedding_dim=base_embedding_dim, name="movie_embedding", )(movie_input) # Compute dot product similarity between user and movie embeddings. logits = layers.Dot(axes=1, name="dot_similarity")( [user_embedding, movie_embedding] ) # Convert to rating scale. prediction = keras.activations.sigmoid(logits) * 5 # Create the model. model = keras.Model( inputs=[user_input, movie_input], outputs=prediction, name="baseline_model" ) return model memory_efficient_model = create_memory_efficient_model() memory_efficient_model.summary() 

history = run_experiment(memory_efficient_model) plt.plot(history.history["loss"]) plt.plot(history.history["val_loss"]) plt.title("model loss") plt.ylabel("loss") plt.xlabel("epoch") plt.legend(["train", "eval"], loc="upper left") plt.show()