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

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.

Start with the linked TensorFlow Recommenders multitask tutorial, then read the shown train_step and compute_loss implementations to understand how rating_weight and retrieval_weight affect the combined objective. Done means documenting whether differing loss magnitudes imply task imbalance and how users should choose or interpret the weights.