Matrix Factorisation with Embeddings

Predict ratings with a matrix factorisation model built from nn.Embedding layers.

Each user and each item gets a vector of size dim and one bias number. The predicted rating is their dot product plus the biases:

r_hat(u, i) = p_u . q_i + b_u + b_i + mu

Write MatrixFactorization(n_users, n_items, dim) with these attributes:

• self.user_emb and self.item_emb: nn.Embedding(n, dim)
• self.user_bias and self.item_bias: nn.Embedding(n, 1)
• self.global_bias: an nn.Parameter scalar mu

forward(users, items) takes two (B,) id tensors and returns (B,) predictions.

Then write train_mf(users, items, ratings, n_users, n_items). It trains a new model with mini-batches from a DataLoader on mean squared error and returns it in eval mode.

Function signature

class MatrixFactorization(torch.nn.Module):
    def __init__(self, n_users, n_items, dim): ...
    def forward(self, users, items): ...


def train_mf(users, items, ratings, n_users, n_items):

Example

Input

p_u = [1.0, 2.0],  q_i = [0.5, -1.0]
b_u = 0.2,  b_i = -0.3,  mu = 3.5

Output

r_hat = 1.9

Reasoning

The dot product is 1.0 * 0.5 + 2.0 * -1.0 = -1.5. Adding the biases gives -1.5 + 0.2 - 0.3 + 3.5 = 1.9. The biases capture "this user rates high" and "this item is popular". The dot product captures taste.

Constraints and notes

• The ratings are synthetic, from a low-rank model plus noise with standard deviation 0.3. About 20% of all user and item pairs are observed.
• The tests need a root mean squared error below 0.6 on held-out pairs. A model with only biases gets about 1.3, so the embeddings must learn.
• With few ratings per user, the model can overfit. A small weight_decay and a modest dim help. Starting mu at the mean training rating speeds things up.