[Question] Weights for retrieval and ranking tasks in multi-task recommenders

[fuchami]

Multi-task Recommendations is a great tutorial! https://www.tensorflow.org/recommenders/examples/multitask

I tried to override train_step to measure rating_loss and retrieval_loss respectively, as shown below.

class MovielensModel(tfrs.models.Model):
  ...

  def train_step(self, inputs):
    with tf.GradientTape() as tape:
      rating_loss, retrieval_loss = self.compute_loss(inputs, training=True)

      total_loss = rating_loss + retrieval_loss

      gradients = tape.gradient(total_loss, self.trainable_variables)
      self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

      metrics = {metric.name: metric.result() for metric in self.metrics}
      metrics['rating_loss'] = rating_loss
      metrics['retrieval_loss'] = retrieval_loss
      metrics['total_loss'] = total_loss

      return metrics

  def compute_loss(self, features: Dict[Text, tf.Tensor], training=False) -> tf.Tensor:

    ratings = features.pop("user_rating")

    user_embeddings, movie_embeddings, rating_predictions = self(features)

    # We compute the loss for each task.
    rating_loss = self.rating_task(
        labels=ratings,
        predictions=rating_predictions,
    )
    retrieval_loss = self.retrieval_task(user_embeddings, movie_embeddings)

    # And combine them using the loss weights.
    return self.rating_weight * rating_loss, self.retrieval_weight * retrieval_loss

train model.

model = MovielensModel(rating_weight=1.0, retrieval_weight=1.0)
model.compile(optimizer=tf.keras.optimizers.Adagrad(0.1))

model.fit(cached_train, epochs=3)

Epoch 1/3
10/10 [==============================] - 16s 2s/step - root_mean_squared_error: 2.4214 - factorized_top_k/top_1_categorical_accuracy: 3.8750e-04 - factorized_top_k/top_5_categorical_accuracy: 0.0044 - factorized_top_k/top_10_categorical_accuracy: 0.0127 - factorized_top_k/top_50_categorical_accuracy: 0.0828 - factorized_top_k/top_100_categorical_accuracy: 0.1602 - rating_loss: 5.3361 - retrieval_loss: 69827.0220 - total_loss: 69832.3580
Epoch 2/3
10/10 [==============================] - 15s 2s/step - root_mean_squared_error: 1.2132 - factorized_top_k/top_1_categorical_accuracy: 0.0026 - factorized_top_k/top_5_categorical_accuracy: 0.0174 - factorized_top_k/top_10_categorical_accuracy: 0.0352 - factorized_top_k/top_50_categorical_accuracy: 0.1624 - factorized_top_k/top_100_categorical_accuracy: 0.2898 - rating_loss: 1.4526 - retrieval_loss: 67463.0369 - total_loss: 67464.4886
Epoch 3/3
10/10 [==============================] - 15s 2s/step - root_mean_squared_error: 1.1004 - factorized_top_k/top_1_categorical_accuracy: 0.0030 - factorized_top_k/top_5_categorical_accuracy: 0.0215 - factorized_top_k/top_10_categorical_accuracy: 0.0437 - factorized_top_k/top_50_categorical_accuracy: 0.1851 - factorized_top_k/top_100_categorical_accuracy: 0.3133 - rating_loss: 1.2014 - retrieval_loss: 66315.0767 - total_loss: 66316.2777

Looking at the model training logs, I see that the scales of loss are different for retrieval and ranking. My question is whether this model fits too well to the retrieval task because the scale of the retrieval loss is larger.

I would greatly appreciate some help from the community from those who are more experienced in the above questions. thank you.