Question: Performance of TFRS vs Keras Models

Hi @maciejkula I am back with more questions :)

I have been working on a non TFRS version of our model for some time, a very simple example of what it looks like is as follows:

inputs = tf.keras.Input(shape=(None,), dtype='int64')
x = embedding_layer(inputs)
x = tf.keras.layers.LSTM(32)(x)
outputs = tf.keras.layers.Dense(128)(x)
model = tf.keras.Model(inputs, outputs)
model.compile(
    optimizer='adam', 
    metrics=[tf.keras.metrics.CosineSimilarity()], 
    loss=tf.keras.losses.CosineSimilarity())
model.fit(X_train, y_train, epochs=30, validation_data=(X_test, y_test))

Where we try and predict the next vector in a sequence using an LSTM layer.

The dataset size lets say is around 300k sequences, with a length of 8-10 items and each item is a vector of 128 dims. The above model takes around 20 secs per epoch.

If I reframe this problem as a recommendation task. The same dataset takes around 5 hours for a single epoch 🤯

My vocabulary of items was large (600k) so I reduced this just to top occurring 10k items and still the change is not huge.

My assumption is that this is due to the Retrieval task, is this correct? Does this task for each sequence in the epoch need to check each candidate embedding?

Hopefully I have done something wrong and there is a possible optimisation, my model code is below:

class Model(tfrs.Model):
    def __init__(self, candidates, user_embedding_layer):

        super().__init__()
        self.candidates = candidates
        self.num_khs = len(candidates)
        self.user_embedding_layer = user_embedding_layer
        self._build_user_model()
        self._build_item_model()
        self._build_task()

    def _build_user_model(self):
        inputs = tf.keras.Input(shape=(None,), dtype='int64')
        x = self.user_embedding_layer(inputs)
        outputs = tf.keras.layers.LSTM(128)(x)
        self.user_model = tf.keras.Model(inputs, outputs)

    def _build_item_model(self):
        inputs = tf.keras.Input(shape=(None,), dtype='string')
        x = tf.keras.layers.experimental.preprocessing.StringLookup(vocabulary=self.candidates)(inputs)
        outputs = tf.keras.layers.Embedding(self.num_khs + 1, 128)(x)
        self.item_model = tf.keras.Model(inputs, outputs)

    def _build_task(self):
        tf_candidates = tf.data.Dataset.from_tensor_slices(self.candidates)
        metrics = tfrs.metrics.FactorizedTopK(
            candidates=tf_candidates.batch(1024).map(self.item_model)
        )
        self.task = tfrs.tasks.Retrieval(
            metrics=metrics
        )

    def compute_loss(self, inputs, training=False):
        sequence, next_item = inputs

        user_emb = self.user_model(sequence)
        next_emb = self.item_model(next_item)

        return self.task(user_emb, next_emb, compute_metrics=not training)


model = Model(vocab, embedding_layer)
model.compile(optimizer='adam')
model.fit(X_train, y_train, epochs=30, validation_data=(X_test, y_test))