From bc6e304fc96402904ed49f8ab583ed4659e5d24c Mon Sep 17 00:00:00 2001 From: danhdanhtuan0308 Date: Sat, 9 May 2026 15:28:49 -0400 Subject: [PATCH 1/3] feat(flexmf): add Neural Collaborative Filtering (NCF) model --- src/lenskit/flexmf/_ncf.py | 146 ++++++++++++++++++++++++++++++++ tests/flexmf/test_flexmf_ncf.py | 48 +++++++++++ 2 files changed, 194 insertions(+) create mode 100644 src/lenskit/flexmf/_ncf.py create mode 100644 tests/flexmf/test_flexmf_ncf.py diff --git a/src/lenskit/flexmf/_ncf.py b/src/lenskit/flexmf/_ncf.py new file mode 100644 index 000000000..4bd04afbe --- /dev/null +++ b/src/lenskit/flexmf/_ncf.py @@ -0,0 +1,146 @@ +# This file is part of LensKit. +# Copyright (C) 2018-2023 Boise State University. +# Copyright (C) 2023-2026 Drexel University. +# Licensed under the MIT license, see LICENSE.md for details. +# SPDX-License-Identifier: MIT + +from __future__ import annotations + +from typing import Literal + +import torch +from pydantic import PositiveInt +from torch import Tensor, nn +from torch.linalg import norm + +from lenskit.logging import get_logger + +from ._base import FlexMFConfigBase, FlexMFScorerBase +from ._model import FlexMFModel +from ._implicit import ImplicitLoss, NegativeStrategy, FlexMFImplicitTrainer + +_log = get_logger(__name__) + + +class FlexMFNCFConfig(FlexMFConfigBase): + """ + Configuration for NCF (Neural Collaborative Filtering) Scorer. + + Stability: + Experimental + """ + gmf_embedding_size: PositiveInt = 8 + mlp_embedding_size: PositiveInt = 8 + mlp_layers: list[PositiveInt] = [16, 8, 4] + + loss: ImplicitLoss = "logistic" + negative_strategy: NegativeStrategy | None = None + negative_count: PositiveInt = 4 + positive_weight: float = 1.0 + + user_bias: bool | None = False + item_bias: bool = False + convolution_layers: int = 0 + + def selected_negative_strategy(self) -> NegativeStrategy: + if self.negative_strategy is not None: + return self.negative_strategy + elif self.loss == "warp": + return "misranked" + else: + return "uniform" + + +class FlexMFNCFModel(FlexMFModel): + """ + Torch module for Neural Collaborative Filtering (NCF). + """ + + def __init__(self, gmf_e_size: int, mlp_e_size: int, mlp_layers: list[int], + n_users: int, n_items: int, rng: torch.Generator, init_scale: float = 0.1, sparse: bool = False): + # We call the super class with layers=0 and without bias. + # This gives us u_embed and i_embed which we use as GMF embeddings. + super().__init__( + e_size=gmf_e_size, n_users=n_users, n_items=n_items, rng=rng, + user_bias=False, item_bias=False, init_scale=init_scale, sparse=sparse, layers=0 + ) + + self.mlp_e_size = mlp_e_size + self.u_mlp_embed = nn.Embedding(n_users, mlp_e_size, sparse=sparse) + self.i_mlp_embed = nn.Embedding(n_items, mlp_e_size, sparse=sparse) + + nn.init.normal_(self.u_mlp_embed.weight, std=init_scale, generator=rng) + nn.init.normal_(self.i_mlp_embed.weight, std=init_scale, generator=rng) + + # Build MLP + mlp_modules = [] + input_size = mlp_e_size * 2 + for size in mlp_layers: + mlp_modules.append(nn.Linear(input_size, size)) + mlp_modules.append(nn.ReLU()) + input_size = size + + self.mlp = nn.Sequential(*mlp_modules) + + # Final output layer + self.prediction = nn.Linear(input_size + gmf_e_size, 1) + nn.init.kaiming_uniform_(self.prediction.weight, a=1, nonlinearity='sigmoid') + + def forward(self, user: Tensor, item: Tensor, *, return_norm: bool = False): + u_gmf = self.u_embed(user) + i_gmf = self.i_embed(item) + + u_mlp = self.u_mlp_embed(user) + i_mlp = self.i_mlp_embed(item) + + # Ensure MLP inputs are broadcasted to the same shape before concatenating + if u_mlp.shape != i_mlp.shape: + u_mlp = u_mlp.expand(i_mlp.shape[:-1] + (u_mlp.shape[-1],)) + i_mlp = i_mlp.expand(u_mlp.shape[:-1] + (i_mlp.shape[-1],)) + + gmf_out = u_gmf * i_gmf + mlp_out = self.mlp(torch.cat([u_mlp, i_mlp], dim=-1)) + + out = torch.cat([gmf_out, mlp_out], dim=-1) + score = self.prediction(out).squeeze(-1) + + if return_norm: + # Return regularizations + l2 = norm(u_gmf, dim=-1) + norm(i_gmf, dim=-1) + norm(u_mlp, dim=-1) + norm(i_mlp, dim=-1) + if l2.shape != score.shape: + # Shape adjustment handled externally if broadcast differs + pass + return torch.stack((score, l2)) + + return score + + +class FlexMFNCFTrainer(FlexMFImplicitTrainer): + """ + Trainer for the NCF Model. Repurposes ImplicitTrainer's loop. + """ + + def create_model(self) -> FlexMFNCFModel: + return FlexMFNCFModel( + gmf_e_size=self.config.gmf_embedding_size, # type: ignore + mlp_e_size=self.config.mlp_embedding_size, # type: ignore + mlp_layers=self.config.mlp_layers, # type: ignore + n_users=self.data.n_users, + n_items=self.data.n_items, + rng=self.torch_rng, + sparse=self.config.reg_method != "AdamW", + ) + + +class FlexMFNCFScorer(FlexMFScorerBase): + """ + Neural Collaborative Filtering (NCF) with FlexMF. + + Stability: + Experimental + """ + + config: FlexMFNCFConfig + + def create_trainer(self, data, options): + return FlexMFNCFTrainer(self, data, options) diff --git a/tests/flexmf/test_flexmf_ncf.py b/tests/flexmf/test_flexmf_ncf.py new file mode 100644 index 000000000..239723360 --- /dev/null +++ b/tests/flexmf/test_flexmf_ncf.py @@ -0,0 +1,48 @@ +# This file is part of LensKit. +# Copyright (C) 2018-2023 Boise State University. +# Copyright (C) 2023-2026 Drexel University. +# Licensed under the MIT license, see LICENSE.md for details. +# SPDX-License-Identifier: MIT + +from itertools import product + +from pytest import mark, skip + +from lenskit.flexmf._ncf import FlexMFNCFConfig, FlexMFNCFScorer +from lenskit.testing import BasicComponentTests, ScorerTests + + +class TestFlexMFNCF(BasicComponentTests, ScorerTests): + expected_ndcg = (0.01, 0.25) + component = FlexMFNCFScorer + config = FlexMFNCFConfig(epochs=3, gmf_embedding_size=16, mlp_embedding_size=16, mlp_layers=[32, 16, 8]) + + def test_skip_retrain(self, ml_ds): + skip("not needed") + + def test_run_with_doubles(self, ml_ratings): + skip("FlexMF is fine with doubles") + + +def test_ncf_config_defaults(): + cfg = FlexMFNCFConfig() + assert cfg.gmf_embedding_size == 8 + assert cfg.mlp_embedding_size == 8 + assert cfg.mlp_layers == [16, 8, 4] + + +def test_ncf_config_negative_default(): + cfg = FlexMFNCFConfig(loss="pairwise") + assert cfg.loss == "pairwise" + assert cfg.selected_negative_strategy() == "uniform" + + +@mark.slow +@mark.parametrize(["loss", "reg"], product(["logistic", "pairwise"], ["AdamW"])) +def test_flexmf_ncf_train_config(ml_ds, loss, reg): + config = FlexMFNCFConfig(loss=loss, reg_method=reg, epochs=1) + model = FlexMFNCFScorer(config) + print("training", model) + model.train(ml_ds) + + assert model.model is not None From ecf0fc7970eb359b5a601c44830036e35cbd1064 Mon Sep 17 00:00:00 2001 From: danhdanhtuan0308 Date: Sat, 9 May 2026 15:47:23 -0400 Subject: [PATCH 2/3] style: fix linting and formatting issues --- src/lenskit/flexmf/_ncf.py | 70 +++++++++++++++++++++------------ tests/flexmf/test_flexmf_ncf.py | 4 +- 2 files changed, 48 insertions(+), 26 deletions(-) diff --git a/src/lenskit/flexmf/_ncf.py b/src/lenskit/flexmf/_ncf.py index 4bd04afbe..49435bc61 100644 --- a/src/lenskit/flexmf/_ncf.py +++ b/src/lenskit/flexmf/_ncf.py @@ -6,8 +6,6 @@ from __future__ import annotations -from typing import Literal - import torch from pydantic import PositiveInt from torch import Tensor, nn @@ -16,8 +14,8 @@ from lenskit.logging import get_logger from ._base import FlexMFConfigBase, FlexMFScorerBase +from ._implicit import FlexMFImplicitTrainer, ImplicitLoss, NegativeStrategy from ._model import FlexMFModel -from ._implicit import ImplicitLoss, NegativeStrategy, FlexMFImplicitTrainer _log = get_logger(__name__) @@ -25,19 +23,20 @@ class FlexMFNCFConfig(FlexMFConfigBase): """ Configuration for NCF (Neural Collaborative Filtering) Scorer. - + Stability: Experimental """ + gmf_embedding_size: PositiveInt = 8 mlp_embedding_size: PositiveInt = 8 mlp_layers: list[PositiveInt] = [16, 8, 4] - + loss: ImplicitLoss = "logistic" negative_strategy: NegativeStrategy | None = None negative_count: PositiveInt = 4 positive_weight: float = 1.0 - + user_bias: bool | None = False item_bias: bool = False convolution_layers: int = 0 @@ -56,22 +55,38 @@ class FlexMFNCFModel(FlexMFModel): Torch module for Neural Collaborative Filtering (NCF). """ - def __init__(self, gmf_e_size: int, mlp_e_size: int, mlp_layers: list[int], - n_users: int, n_items: int, rng: torch.Generator, init_scale: float = 0.1, sparse: bool = False): + def __init__( + self, + gmf_e_size: int, + mlp_e_size: int, + mlp_layers: list[int], + n_users: int, + n_items: int, + rng: torch.Generator, + init_scale: float = 0.1, + sparse: bool = False, + ): # We call the super class with layers=0 and without bias. # This gives us u_embed and i_embed which we use as GMF embeddings. super().__init__( - e_size=gmf_e_size, n_users=n_users, n_items=n_items, rng=rng, - user_bias=False, item_bias=False, init_scale=init_scale, sparse=sparse, layers=0 + e_size=gmf_e_size, + n_users=n_users, + n_items=n_items, + rng=rng, + user_bias=False, + item_bias=False, + init_scale=init_scale, + sparse=sparse, + layers=0, ) - + self.mlp_e_size = mlp_e_size self.u_mlp_embed = nn.Embedding(n_users, mlp_e_size, sparse=sparse) self.i_mlp_embed = nn.Embedding(n_items, mlp_e_size, sparse=sparse) - + nn.init.normal_(self.u_mlp_embed.weight, std=init_scale, generator=rng) nn.init.normal_(self.i_mlp_embed.weight, std=init_scale, generator=rng) - + # Build MLP mlp_modules = [] input_size = mlp_e_size * 2 @@ -79,39 +94,44 @@ def __init__(self, gmf_e_size: int, mlp_e_size: int, mlp_layers: list[int], mlp_modules.append(nn.Linear(input_size, size)) mlp_modules.append(nn.ReLU()) input_size = size - + self.mlp = nn.Sequential(*mlp_modules) - + # Final output layer self.prediction = nn.Linear(input_size + gmf_e_size, 1) - nn.init.kaiming_uniform_(self.prediction.weight, a=1, nonlinearity='sigmoid') + nn.init.kaiming_uniform_(self.prediction.weight, a=1, nonlinearity="sigmoid") def forward(self, user: Tensor, item: Tensor, *, return_norm: bool = False): u_gmf = self.u_embed(user) i_gmf = self.i_embed(item) - + u_mlp = self.u_mlp_embed(user) i_mlp = self.i_mlp_embed(item) - + # Ensure MLP inputs are broadcasted to the same shape before concatenating if u_mlp.shape != i_mlp.shape: u_mlp = u_mlp.expand(i_mlp.shape[:-1] + (u_mlp.shape[-1],)) i_mlp = i_mlp.expand(u_mlp.shape[:-1] + (i_mlp.shape[-1],)) - + gmf_out = u_gmf * i_gmf mlp_out = self.mlp(torch.cat([u_mlp, i_mlp], dim=-1)) - + out = torch.cat([gmf_out, mlp_out], dim=-1) score = self.prediction(out).squeeze(-1) - + if return_norm: # Return regularizations - l2 = norm(u_gmf, dim=-1) + norm(i_gmf, dim=-1) + norm(u_mlp, dim=-1) + norm(i_mlp, dim=-1) + l2 = ( + norm(u_gmf, dim=-1) + + norm(i_gmf, dim=-1) + + norm(u_mlp, dim=-1) + + norm(i_mlp, dim=-1) + ) if l2.shape != score.shape: # Shape adjustment handled externally if broadcast differs pass return torch.stack((score, l2)) - + return score @@ -119,7 +139,7 @@ class FlexMFNCFTrainer(FlexMFImplicitTrainer): """ Trainer for the NCF Model. Repurposes ImplicitTrainer's loop. """ - + def create_model(self) -> FlexMFNCFModel: return FlexMFNCFModel( gmf_e_size=self.config.gmf_embedding_size, # type: ignore @@ -135,7 +155,7 @@ def create_model(self) -> FlexMFNCFModel: class FlexMFNCFScorer(FlexMFScorerBase): """ Neural Collaborative Filtering (NCF) with FlexMF. - + Stability: Experimental """ diff --git a/tests/flexmf/test_flexmf_ncf.py b/tests/flexmf/test_flexmf_ncf.py index 239723360..8ca313879 100644 --- a/tests/flexmf/test_flexmf_ncf.py +++ b/tests/flexmf/test_flexmf_ncf.py @@ -15,7 +15,9 @@ class TestFlexMFNCF(BasicComponentTests, ScorerTests): expected_ndcg = (0.01, 0.25) component = FlexMFNCFScorer - config = FlexMFNCFConfig(epochs=3, gmf_embedding_size=16, mlp_embedding_size=16, mlp_layers=[32, 16, 8]) + config = FlexMFNCFConfig( + epochs=3, gmf_embedding_size=16, mlp_embedding_size=16, mlp_layers=[32, 16, 8] + ) def test_skip_retrain(self, ml_ds): skip("not needed") From 9bd453094806d1c90b42d1c2793ae76a98731d49 Mon Sep 17 00:00:00 2001 From: danhdanhtuan0308 Date: Sun, 24 May 2026 13:44:39 -0400 Subject: [PATCH 3/3] refactor(flexmf): address NCF PR review comments --- src/lenskit/flexmf/__init__.py | 3 ++ src/lenskit/flexmf/_ncf.py | 81 ++++++++++++++++++++------------- tests/flexmf/test_flexmf_ncf.py | 4 +- 3 files changed, 55 insertions(+), 33 deletions(-) diff --git a/src/lenskit/flexmf/__init__.py b/src/lenskit/flexmf/__init__.py index 171610d7e..38ce52430 100644 --- a/src/lenskit/flexmf/__init__.py +++ b/src/lenskit/flexmf/__init__.py @@ -20,6 +20,7 @@ from ._base import FlexMFConfigBase, FlexMFScorerBase from ._explicit import FlexMFExplicitConfig, FlexMFExplicitScorer from ._implicit import FlexMFImplicitConfig, FlexMFImplicitScorer +from ._ncf import FlexMFNCFConfig, FlexMFNCFScorer __all__ = [ "FlexMFConfigBase", @@ -28,4 +29,6 @@ "FlexMFExplicitScorer", "FlexMFImplicitConfig", "FlexMFImplicitScorer", + "FlexMFNCFConfig", + "FlexMFNCFScorer", ] diff --git a/src/lenskit/flexmf/_ncf.py b/src/lenskit/flexmf/_ncf.py index 49435bc61..665e3f525 100644 --- a/src/lenskit/flexmf/_ncf.py +++ b/src/lenskit/flexmf/_ncf.py @@ -13,48 +13,51 @@ from lenskit.logging import get_logger -from ._base import FlexMFConfigBase, FlexMFScorerBase -from ._implicit import FlexMFImplicitTrainer, ImplicitLoss, NegativeStrategy +from ._base import FlexMFScorerBase +from ._implicit import FlexMFImplicitConfig, FlexMFImplicitTrainer from ._model import FlexMFModel _log = get_logger(__name__) -class FlexMFNCFConfig(FlexMFConfigBase): +class FlexMFNCFConfig(FlexMFImplicitConfig): """ - Configuration for NCF (Neural Collaborative Filtering) Scorer. + Configuration for NCF (Neural Collaborative Filtering) Scorer. It inherits + common training options and implicit-feedback settings from + :class:`FlexMFImplicitConfig`. The inherited ``embedding_size`` field is + used as the GMF embedding size. Stability: Experimental """ - gmf_embedding_size: PositiveInt = 8 mlp_embedding_size: PositiveInt = 8 + """ + The size of the MLP embedding space. + """ + mlp_layers: list[PositiveInt] = [16, 8, 4] + """ + The sizes of the MLP hidden layers. + """ - loss: ImplicitLoss = "logistic" - negative_strategy: NegativeStrategy | None = None + # Override implicit defaults: NCF uses more negatives and no bias terms by default. negative_count: PositiveInt = 4 - positive_weight: float = 1.0 - user_bias: bool | None = False item_bias: bool = False - convolution_layers: int = 0 - - def selected_negative_strategy(self) -> NegativeStrategy: - if self.negative_strategy is not None: - return self.negative_strategy - elif self.loss == "warp": - return "misranked" - else: - return "uniform" -class FlexMFNCFModel(FlexMFModel): +class FlexMFNCFModel(nn.Module): """ Torch module for Neural Collaborative Filtering (NCF). + + Uses composition rather than inheritance: holds a :class:`FlexMFModel` for + the GMF path alongside a separate MLP tower, combining both via a final + linear layer. """ + gmf_model: FlexMFModel + def __init__( self, gmf_e_size: int, @@ -66,9 +69,10 @@ def __init__( init_scale: float = 0.1, sparse: bool = False, ): - # We call the super class with layers=0 and without bias. - # This gives us u_embed and i_embed which we use as GMF embeddings. - super().__init__( + super().__init__() + + # GMF component: a standard MF model providing u_embed / i_embed. + self.gmf_model = FlexMFModel( e_size=gmf_e_size, n_users=n_users, n_items=n_items, @@ -97,39 +101,54 @@ def __init__( self.mlp = nn.Sequential(*mlp_modules) - # Final output layer + # Final output layer: concatenated GMF and MLP outputs → scalar score. self.prediction = nn.Linear(input_size + gmf_e_size, 1) nn.init.kaiming_uniform_(self.prediction.weight, a=1, nonlinearity="sigmoid") + @property + def device(self): + return self.gmf_model.device + + def zero_users(self, users: Tensor): + """Zero weights for users with no training interactions.""" + self.gmf_model.zero_users(users) + self.u_mlp_embed.weight.data[users] = 0 + + def zero_items(self, items: Tensor): + """Zero weights for items with no training interactions.""" + self.gmf_model.zero_items(items) + self.i_mlp_embed.weight.data[items] = 0 + def forward(self, user: Tensor, item: Tensor, *, return_norm: bool = False): - u_gmf = self.u_embed(user) - i_gmf = self.i_embed(item) + u_gmf = self.gmf_model.u_embed(user) + i_gmf = self.gmf_model.i_embed(item) u_mlp = self.u_mlp_embed(user) i_mlp = self.i_mlp_embed(item) - # Ensure MLP inputs are broadcasted to the same shape before concatenating + # Broadcast MLP embeddings to the same shape before concatenating. if u_mlp.shape != i_mlp.shape: u_mlp = u_mlp.expand(i_mlp.shape[:-1] + (u_mlp.shape[-1],)) i_mlp = i_mlp.expand(u_mlp.shape[:-1] + (i_mlp.shape[-1],)) + # GMF path: element-wise product of user and item embeddings. gmf_out = u_gmf * i_gmf + # MLP path: concatenate embeddings and pass through the MLP tower. mlp_out = self.mlp(torch.cat([u_mlp, i_mlp], dim=-1)) + # Combine GMF and MLP outputs, then project to a scalar. + # squeeze(-1) removes the trailing size-1 dimension produced by the + # linear layer, giving shape (...) instead of (..., 1). out = torch.cat([gmf_out, mlp_out], dim=-1) score = self.prediction(out).squeeze(-1) if return_norm: - # Return regularizations l2 = ( norm(u_gmf, dim=-1) + norm(i_gmf, dim=-1) + norm(u_mlp, dim=-1) + norm(i_mlp, dim=-1) ) - if l2.shape != score.shape: - # Shape adjustment handled externally if broadcast differs - pass return torch.stack((score, l2)) return score @@ -142,7 +161,7 @@ class FlexMFNCFTrainer(FlexMFImplicitTrainer): def create_model(self) -> FlexMFNCFModel: return FlexMFNCFModel( - gmf_e_size=self.config.gmf_embedding_size, # type: ignore + gmf_e_size=self.config.embedding_size, mlp_e_size=self.config.mlp_embedding_size, # type: ignore mlp_layers=self.config.mlp_layers, # type: ignore n_users=self.data.n_users, diff --git a/tests/flexmf/test_flexmf_ncf.py b/tests/flexmf/test_flexmf_ncf.py index 8ca313879..b9bc90028 100644 --- a/tests/flexmf/test_flexmf_ncf.py +++ b/tests/flexmf/test_flexmf_ncf.py @@ -16,7 +16,7 @@ class TestFlexMFNCF(BasicComponentTests, ScorerTests): expected_ndcg = (0.01, 0.25) component = FlexMFNCFScorer config = FlexMFNCFConfig( - epochs=3, gmf_embedding_size=16, mlp_embedding_size=16, mlp_layers=[32, 16, 8] + epochs=3, embedding_size=16, mlp_embedding_size=16, mlp_layers=[32, 16, 8] ) def test_skip_retrain(self, ml_ds): @@ -28,7 +28,7 @@ def test_run_with_doubles(self, ml_ratings): def test_ncf_config_defaults(): cfg = FlexMFNCFConfig() - assert cfg.gmf_embedding_size == 8 + assert cfg.embedding_size == 64 # inherited default from FlexMFConfigBase assert cfg.mlp_embedding_size == 8 assert cfg.mlp_layers == [16, 8, 4]