本記事は ZOZO Advent Calendar 2023 シリーズ 7 の 1 日目の記事です。

概要

本記事では Two-Tower モデルを使って user2item 推薦モデルを構築した際に得られる embeddings を活用し、簡単に item2item, user2user, item2user 推薦を実現する方法についてお伝えします。この方法は TensorFlow Recommenders プロジェクトの Issue 内の議論から着想を得ました。

まずは上記で説明した 4 種類の推薦と用途について以下で説明します。
user2item, item2item は広く知られる推薦手法である一方、user2user, item2user 推薦は用途が限定的だったため考えてみました。

推薦の種類	用途
user2item	ユーザーに関連度の高いアイテムを推薦
item2item	ユーザーに特定のアイテムの類似アイテムを推薦アイテムセグメントの抽出
user2user	ユーザーに類似のインフルエンサーなどを推薦ユーザーセグメントの抽出
item2user	在庫が限られるアイテムをユーザーに推薦

実装の前に、今回は user2item 推薦のタスクを解く中で item2item, user2user, item2user 推薦も可能になるという話ですが、それぞれの推薦をメインタスクとした学習方法が本筋であることを補足しておきます。

モデルの実装

以下では TensorFlow Recommenders チュートリアル > Building deep retrieval models に基づいて Two-Tower モデルの実装を行います。
学習データは、ユーザーの映画に対するレビューデータが含まれる MovieLens を使用し TensorFlow Recommenders を使用してユーザーの特徴量から映画を推薦するタスクを実装します。

以下のサンプルコードは Jupyter Notebook 上で動かしてください。

サンプルコード

import os
import tempfile

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_recommenders as tfrs

plt.style.use("seaborn-whitegrid")

ratings = tfds.load("movielens/100k-ratings", split="train")
movies = tfds.load("movielens/100k-movies", split="train")

ratings = ratings.map(
    lambda x: {
        "movie_title": x["movie_title"],
        "user_id": x["user_id"],
        "timestamp": x["timestamp"],
    }
)
movies = movies.map(lambda x: x["movie_title"])

timestamps = np.concatenate(
    list(ratings.map(lambda x: x["timestamp"]).batch(100))
)

max_timestamp = timestamps.max()
min_timestamp = timestamps.min()

timestamp_buckets = np.linspace(
    min_timestamp,
    max_timestamp,
    num=1000,
)

unique_movie_titles = np.unique(np.concatenate(list(movies.batch(1000))))
unique_user_ids = np.unique(
    np.concatenate(list(ratings.batch(1_000).map(lambda x: x["user_id"])))
)


class UserModel(tf.keras.Model):
    def __init__(self):
        super().__init__()

        self.user_embedding = tf.keras.Sequential(
            [
                tf.keras.layers.StringLookup(
                    vocabulary=unique_user_ids, mask_token=None
                ),
                tf.keras.layers.Embedding(len(unique_user_ids) + 1, 32),
            ]
        )
        self.timestamp_embedding = tf.keras.Sequential(
            [
                tf.keras.layers.Discretization(timestamp_buckets.tolist()),
                tf.keras.layers.Embedding(len(timestamp_buckets) + 1, 32),
            ]
        )
        self.normalized_timestamp = tf.keras.layers.Normalization(axis=None)

        self.normalized_timestamp.adapt(timestamps)

    def call(self, inputs):
        # Take the input dictionary, pass it through each input layer,
        # and concatenate the result.
        return tf.concat(
            [
                self.user_embedding(inputs["user_id"]),
                self.timestamp_embedding(inputs["timestamp"]),
                tf.reshape(
                    self.normalized_timestamp(inputs["timestamp"]), (-1, 1)
                ),
            ],
            axis=1,
        )


class QueryModel(tf.keras.Model):
    """Model for encoding user queries."""

    def __init__(self, layer_sizes):
        """Model for encoding user queries.

        Args:
          layer_sizes:
            A list of integers where the i-th entry represents the number of units
            the i-th layer contains.
        """
        super().__init__()

        # We first use the user model for generating embeddings.
        self.embedding_model = UserModel()

        # Then construct the layers.
        self.dense_layers = tf.keras.Sequential()

        # Use the ReLU activation for all but the last layer.
        for layer_size in layer_sizes[:-1]:
            self.dense_layers.add(
                tf.keras.layers.Dense(layer_size, activation="relu")
            )

        # No activation for the last layer.
        for layer_size in layer_sizes[-1:]:
            self.dense_layers.add(tf.keras.layers.Dense(layer_size))

    def call(self, inputs):
        feature_embedding = self.embedding_model(inputs)
        return self.dense_layers(feature_embedding)


class MovieModel(tf.keras.Model):
    def __init__(self):
        super().__init__()

        max_tokens = 10_000

        self.title_embedding = tf.keras.Sequential(
            [
                tf.keras.layers.StringLookup(
                    vocabulary=unique_movie_titles, mask_token=None
                ),
                tf.keras.layers.Embedding(len(unique_movie_titles) + 1, 32),
            ]
        )

        self.title_vectorizer = tf.keras.layers.TextVectorization(
            max_tokens=max_tokens
        )

        self.title_text_embedding = tf.keras.Sequential(
            [
                self.title_vectorizer,
                tf.keras.layers.Embedding(max_tokens, 32, mask_zero=True),
                tf.keras.layers.GlobalAveragePooling1D(),
            ]
        )

        self.title_vectorizer.adapt(movies)

    def call(self, titles):
        return tf.concat(
            [
                self.title_embedding(titles),
                self.title_text_embedding(titles),
            ],
            axis=1,
        )


class CandidateModel(tf.keras.Model):
    """Model for encoding movies."""

    def __init__(self, layer_sizes):
        """Model for encoding movies.

        Args:
          layer_sizes:
            A list of integers where the i-th entry represents the number of units
            the i-th layer contains.
        """
        super().__init__()

        self.embedding_model = MovieModel()

        # Then construct the layers.
        self.dense_layers = tf.keras.Sequential()

        # Use the ReLU activation for all but the last layer.
        for layer_size in layer_sizes[:-1]:
            self.dense_layers.add(
                tf.keras.layers.Dense(layer_size, activation="relu")
            )

        # No activation for the last layer.
        for layer_size in layer_sizes[-1:]:
            self.dense_layers.add(tf.keras.layers.Dense(layer_size))

    def call(self, inputs):
        feature_embedding = self.embedding_model(inputs)
        return self.dense_layers(feature_embedding)


class MovielensModel(tfrs.models.Model):
    def __init__(self, layer_sizes):
        super().__init__()
        self.query_model = QueryModel(layer_sizes)
        self.candidate_model = CandidateModel(layer_sizes)
        self.task = tfrs.tasks.Retrieval(
            metrics=tfrs.metrics.FactorizedTopK(
                candidates=movies.batch(128).map(self.candidate_model),
            ),
        )

    def compute_loss(self, features, training=False):
        # We only pass the user id and timestamp features into the query model. This
        # is to ensure that the training inputs would have the same keys as the
        # query inputs. Otherwise the discrepancy in input structure would cause an
        # error when loading the query model after saving it.
        query_embeddings = self.query_model(
            {
                "user_id": features["user_id"],
                "timestamp": features["timestamp"],
            }
        )
        movie_embeddings = self.candidate_model(features["movie_title"])

        return self.task(
            query_embeddings, movie_embeddings, compute_metrics=not training
        )


tf.random.set_seed(42)
shuffled = ratings.shuffle(100_000, seed=42, reshuffle_each_iteration=False)

train = shuffled.take(80_000)
test = shuffled.skip(80_000).take(20_000)

cached_train = train.shuffle(100_000).batch(2048)
cached_test = test.batch(4096).cache()

num_epochs = 5

model = MovielensModel([32])
model.compile(optimizer=tf.keras.optimizers.Adagrad(0.1))

one_layer_history = model.fit(
    cached_train,
    validation_data=cached_test,
    validation_freq=1,
    epochs=num_epochs,
    verbose=1,
)

accuracy = one_layer_history.history[
    "val_factorized_top_k/top_100_categorical_accuracy"
][-1]
print(f"Top-100 accuracy: {accuracy:.2f}.")

4種類の推薦

上記でモデルの学習が完了したので、以下で冒頭の4種類の推薦の実装方法を説明します。

user2item

user2item は TensorFlow Recommenders チュートリアル > Recommending movies: retrieval を参考にした実装です。
ユーザーの id と timestamp の特徴量から、関連度の高い映画を推薦します。

▼ 実装

# user2item
index = tfrs.layers.factorized_top_k.BruteForce(model.query_model)
index.index_from_dataset(
    tf.data.Dataset.zip(
        (movies.batch(100), movies.batch(100).map(model.candidate_model))
    )
)
_, titles = index(
    {"user_id": tf.constant(["42"]), "timestamp": tf.constant([879024327])}
)
display(titles.numpy()[0].tolist())

▼ output

[b'Foxfire (1996)',
 b'Kazaam (1996)',
 b'Mary Reilly (1996)',
 b'Bed of Roses (1996)',
 b'Zeus and Roxanne (1997)',
 b'Diabolique (1996)',
 b'Phenomenon (1996)',
 b'Homeward Bound II: Lost in San Francisco (1996)',
 b'Juror, The (1996)',
 b'Matilda (1996)']

item2item

item2item はチュートリアルに記載されていませんが、実用的な方法だと思います。
例では、Toy Story (1995) に対して類似度の高い映画を推薦しています。

▼ 実装

# item2item
index = tfrs.layers.factorized_top_k.BruteForce(model.candidate_model)
index.index_from_dataset(
    tf.data.Dataset.zip(
        (movies.batch(100), movies.batch(100).map(model.candidate_model))
    )
)
_, titles = index(tf.constant(["Toy Story (1995)"]))
display(titles.numpy()[0].tolist())

▼ output

[b'Toy Story (1995)',
 b"Mr. Holland's Opus (1995)",
 b'Willy Wonka and the Chocolate Factory (1971)',
 b'Men in Black (1997)',
 b'Hercules (1997)',
 b'Return of the Jedi (1983)',
 b'Rock, The (1996)',
 b'Dead Man Walking (1995)',
 b'Star Wars (1977)',
 b'Ransom (1996)']

use2user

user2user もチュートリアルに記載されていませんが、実現可能です。
あるユーザーの特徴量を入力とし、類似度の高いユーザー(user_id)を推薦します。
今回はモデルの都合上 user_id は同じだが timestamp が違うデータがあるため、推薦上位 500件を取得し重複削除しています。

▼ 実装

# user2user
import pandas as pd

index = tfrs.layers.factorized_top_k.BruteForce(model.query_model)
index.index_from_dataset(
    tf.data.Dataset.zip(
        (
            ratings.batch(100).map(lambda x: x["user_id"]),
            ratings.batch(100)
            .map(
                lambda features: {
                    "user_id": features["user_id"],
                    "timestamp": features["timestamp"],
                }
            )
            .map(model.query_model),
        )
    )
)

_, users = index(
    {"user_id": tf.constant(["42"]), "timestamp": tf.constant([879024327])},
    k=500,
)
display(pd.unique(users.numpy()[0]).tolist())

▼ output

[b'337',
 b'792',
 b'665',
 b'605',
 b'891',
 b'159',
 b'708',
 b'517',
 b'525',
 b'779',
 b'82',
 b'152',
 b'689',
 b'231',
 b'795',
 b'290',
 b'93',
 b'935',
 b'549',
 b'684',
 b'825',
 b'927',
 b'45',
 b'289']

item2user

item2user もチュートリアルに記載されていませんが、実現可能です。
ある映画の特徴量を入力とし、関連度の高いユーザー(user_id)を推薦します。
こちらも推薦上位 500件取得し、重複削除しています。

▼ 実装

# item2user
index = tfrs.layers.factorized_top_k.BruteForce(model.candidate_model)
index.index_from_dataset(
    tf.data.Dataset.zip(
        (
            ratings.batch(100).map(lambda x: x["user_id"]),
            ratings.batch(100)
            .map(
                lambda features: {
                    "user_id": features["user_id"],
                    "timestamp": features["timestamp"],
                }
            )
            .map(model.query_model),
        )
    )
)

_, users = index(
    tf.constant(["Toy Story (1995)"]),
    k=500,
)
display(pd.unique(users.numpy()[0]).tolist())

▼ output

[b'742',
 b'17',
 b'779',
 b'800',
 b'634',
 b'689',
 b'549',
 b'941',
 b'759',
 b'708',
 b'231',
 b'718',
 b'45',
 b'182',
 b'157',
 b'168',
 b'66',
 b'93',
 b'238',
 b'649',
 b'714',
 b'937',
 b'525']

まとめ

本記事では Two-Tower モデルを使って user2item 推薦モデルを構築した際に得られる embeddings を活用し、簡単に item2item, user2user, item2user 推薦を実現する方法についてお伝えしました。

ご紹介した方法は Two-Tower モデル以外の例えば Matrix Factorization でも利用可能です。類似度の計算と近傍探索は別で行う必要がありますが、学習の副産物であるアイテムとユーザーの embeddings を利用することで、item2item, user2user の推薦が可能です。

また、ご紹介した方法は、状況によって実現したい推薦とアンマッチな場合があるため、その際は実現したい推薦をメインタスクとしたモデルを構築することを考えてみてください。