refactor how we initialize for higher scale on implicit-feedback

Author: mdekstrand

src/lenskit/flexmf/_base.py

@@ -141,17 +141,14 @@ def prepare_data(
         """
         raise NotImplementedError()
 
-    def create_model(self, context: FlexMFTrainingContext, data: FlexMFTrainingData) -> FlexMFModel:
+    @abstractmethod
+    def create_model(
+        self, context: FlexMFTrainingContext, data: FlexMFTrainingData
+    ) -> FlexMFModel:  # pragma: nocover
         """
         Prepare the model for training.
         """
-        return FlexMFModel(
-            self.config.embedding_size,
-            data.n_users,
-            data.n_items,
-            context.torch_rng,
-            sparse=self.config.reg_method != "AdamW",
-        )
+        raise NotImplementedError()
 
     def create_optimizer(self, context: FlexMFTrainingContext) -> torch.optim.Optimizer:
         """

src/lenskit/flexmf/_explicit.py

@@ -12,6 +12,7 @@
 from torch.nn import functional as F
 
 from lenskit.data import Dataset
+from lenskit.flexmf._model import FlexMFModel
 from lenskit.training import TrainingOptions
 
 from ._base import FlexMFConfigBase, FlexMFScorerBase
@@ -74,6 +75,19 @@ def prepare_data(
             fields={"ratings": rm_values},
         ).to(context.device)
 
+    def create_model(self, context: FlexMFTrainingContext, data: FlexMFTrainingData) -> FlexMFModel:
+        """
+        Prepare the model for training.
+        """
+        return FlexMFModel(
+            self.config.embedding_size,
+            data.n_users,
+            data.n_items,
+            context.torch_rng,
+            sparse=self.config.reg_method != "AdamW",
+            init_scale=0.1,
+        )
+
     def train_batch(
         self, context: FlexMFTrainingContext, batch: FlexMFTrainingBatch, opt: torch.optim.Optimizer
     ) -> float:

src/lenskit/flexmf/_model.py

@@ -43,6 +43,7 @@ def __init__(
         n_users: int,
         n_items: int,
         rng: torch.Generator,
+        init_scale: float = 1.0,
         user_bias: bool = True,
         item_bias: bool = True,
         sparse: bool = False,
@@ -64,12 +65,12 @@ def __init__(
 
         # initialize all values to a small normal
         if self.u_bias is not None:
-            nn.init.normal_(self.u_bias.weight, std=0.05, generator=rng)
+            nn.init.normal_(self.u_bias.weight, std=init_scale, generator=rng)
         if self.i_bias is not None:
-            nn.init.normal_(self.i_bias.weight, std=0.05, generator=rng)
+            nn.init.normal_(self.i_bias.weight, std=init_scale, generator=rng)
 
-        nn.init.normal_(self.u_embed.weight, std=0.05, generator=rng)
-        nn.init.normal_(self.i_embed.weight, std=0.05, generator=rng)
+        nn.init.normal_(self.u_embed.weight, std=init_scale, generator=rng)
+        nn.init.normal_(self.i_embed.weight, std=init_scale, generator=rng)
 
     @property
     def device(self):