nextrec 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

nextrec/models/ranking/fm.py

@@ -1,15 +1,52 @@
 """
 Date: create on 09/11/2025
-Author:
-    Yang Zhou,zyaztec@gmail.com
+Checkpoint: edit on 09/12/2025
+Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-    [1] Rendle S. Factorization machines[C]//ICDM. 2010: 995-1000.
+[1] Rendle S. Factorization machines[C]//ICDM. 2010: 995-1000.
+
+Factorization Machines (FM) capture second-order feature interactions with
+linear complexity by factorizing the pairwise interaction matrix. Each field
+is embedded into a latent vector; FM models the dot product of every pair of
+embeddings and sums them along with a linear term, enabling strong performance
+with sparse high-dimensional data and minimal feature engineering.
+
+Pipeline:
+  (1) Embed sparse and sequence fields into low-dimensional vectors
+  (2) Compute linear logit over concatenated embeddings
+  (3) Compute pairwise interaction logit via factorized dot products
+  (4) Sum linear + interaction terms and apply prediction layer
+
+Key Advantages:
+- Models pairwise interactions efficiently (O(nk) vs. O(n^2))
+- Works well on sparse inputs without handcrafted crosses
+- Simple architecture with strong baseline performance
+
+FM 是一种通过分解二阶特征交互矩阵、以线性复杂度建模特征对的 CTR 模型。
+每个特征映射为低维向量，FM 对任意特征对进行内积求和并叠加线性项，
+无需复杂特征工程即可在稀疏高维场景取得稳健效果。
+
+处理流程：
+  (1) 对稀疏/序列特征做 embedding
+  (2) 计算线性部分的 logit
+  (3) 计算嵌入对之间的二阶交互 logit
+  (4) 线性项与交互项求和，再通过预测层输出
+
+主要优点：
+- 线性复杂度建模二阶交互，效率高
+- 对稀疏特征友好，减少人工特征交叉
+- 结构简单、表现强健，常作 CTR 基线
 """
 
 import torch.nn as nn
 
 from nextrec.basic.model import BaseModel
-from nextrec.basic.layers import EmbeddingLayer, FM as FMInteraction, LR, PredictionLayer
+from nextrec.basic.layers import (
+    EmbeddingLayer,
+    FM as FMInteraction,
+    LR,
+    PredictionLayer,
+)
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 
 
@@ -21,24 +58,30 @@ class FM(BaseModel):
     @property
     def default_task(self):
         return "binary"
-    
-    def __init__(self,
-                 dense_features: list[DenseFeature] | list = [],
-                 sparse_features: list[SparseFeature] | list = [],
-                 sequence_features: list[SequenceFeature] | list = [],
-                 target: list[str] | list = [],
-                 task: str | list[str] | None = None,
-                 optimizer: str = "adam",
-                 optimizer_params: dict = {},
-                 loss: str | nn.Module | None = "bce",
-                 loss_params: dict | list[dict] | None = None,
-                 device: str = 'cpu',
-                 embedding_l1_reg=1e-6,
-                 dense_l1_reg=1e-5,
-                 embedding_l2_reg=1e-5,
-                 dense_l2_reg=1e-4,
-                 **kwargs):
-        
+
+    def __init__(
+        self,
+        dense_features: list[DenseFeature] | None = None,
+        sparse_features: list[SparseFeature] | None = None,
+        sequence_features: list[SequenceFeature] | None = None,
+        target: list[str] | str | None = None,
+        task: str | list[str] | None = None,
+        optimizer: str = "adam",
+        optimizer_params: dict | None = None,
+        loss: str | nn.Module | None = "bce",
+        loss_params: dict | list[dict] | None = None,
+        device: str = "cpu",
+        embedding_l1_reg=1e-6,
+        dense_l1_reg=1e-5,
+        embedding_l2_reg=1e-5,
+        dense_l2_reg=1e-4,
+        **kwargs,
+    ):
+
+        dense_features = dense_features or []
+        sparse_features = sparse_features or []
+        sequence_features = sequence_features or []
+
         super(FM, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -50,13 +93,13 @@ class FM(BaseModel):
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            **kwargs
+            **kwargs,
         )
 
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-            
+
         self.fm_features = sparse_features + sequence_features
         if len(self.fm_features) == 0:
             raise ValueError("FM requires at least one sparse or sequence feature.")
@@ -70,8 +113,7 @@ class FM(BaseModel):
 
         # Register regularization weights
         self.register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=['linear']
+            embedding_attr="embedding", include_modules=["linear"]
         )
 
         self.compile(

nextrec/models/ranking/masknet.py

@@ -69,12 +69,13 @@ class InstanceGuidedMask(nn.Module):
         self.fc2 = nn.Linear(hidden_dim, output_dim)
 
     def forward(self, v_emb_flat: torch.Tensor) -> torch.Tensor:
-        # v_emb_flat: [batch, features count * embedding_dim] 
+        # v_emb_flat: [batch, features count * embedding_dim]
         x = self.fc1(v_emb_flat)
         x = F.relu(x)
         v_mask = self.fc2(x)
         return v_mask
 
+
 class MaskBlockOnEmbedding(nn.Module):
     def __init__(
         self,
@@ -86,20 +87,28 @@ class MaskBlockOnEmbedding(nn.Module):
         super().__init__()
         self.num_fields = num_fields
         self.embedding_dim = embedding_dim
-        self.input_dim = num_fields * embedding_dim  # input_dim = features count * embedding_dim
+        self.input_dim = (
+            num_fields * embedding_dim
+        )  # input_dim = features count * embedding_dim
         self.ln_emb = nn.LayerNorm(embedding_dim)
-        self.mask_gen = InstanceGuidedMask(input_dim=self.input_dim, hidden_dim=mask_hidden_dim, output_dim=self.input_dim,)
+        self.mask_gen = InstanceGuidedMask(
+            input_dim=self.input_dim,
+            hidden_dim=mask_hidden_dim,
+            output_dim=self.input_dim,
+        )
         self.ffn = nn.Linear(self.input_dim, hidden_dim)
         self.ln_hid = nn.LayerNorm(hidden_dim)
 
     # different from MaskBlockOnHidden: input is field embeddings
-    def forward(self, field_emb: torch.Tensor, v_emb_flat: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, field_emb: torch.Tensor, v_emb_flat: torch.Tensor
+    ) -> torch.Tensor:
         B = field_emb.size(0)
-        norm_emb = self.ln_emb(field_emb)        # [B, features count, embedding_dim]
-        norm_emb_flat = norm_emb.view(B, -1)     # [B, features count * embedding_dim]
-        v_mask = self.mask_gen(v_emb_flat)       # [B, features count * embedding_dim]
-        v_masked_emb = v_mask * norm_emb_flat    # [B, features count * embedding_dim]
-        hidden = self.ffn(v_masked_emb)          # [B, hidden_dim]
+        norm_emb = self.ln_emb(field_emb)  # [B, features count, embedding_dim]
+        norm_emb_flat = norm_emb.view(B, -1)  # [B, features count * embedding_dim]
+        v_mask = self.mask_gen(v_emb_flat)  # [B, features count * embedding_dim]
+        v_masked_emb = v_mask * norm_emb_flat  # [B, features count * embedding_dim]
+        hidden = self.ffn(v_masked_emb)  # [B, hidden_dim]
         hidden = self.ln_hid(hidden)
         hidden = F.relu(hidden)
 
@@ -123,15 +132,21 @@ class MaskBlockOnHidden(nn.Module):
         self.ln_input = nn.LayerNorm(hidden_dim)
         self.ln_output = nn.LayerNorm(hidden_dim)
 
-        self.mask_gen = InstanceGuidedMask(input_dim=self.v_emb_dim, hidden_dim=mask_hidden_dim, output_dim=hidden_dim,)
+        self.mask_gen = InstanceGuidedMask(
+            input_dim=self.v_emb_dim,
+            hidden_dim=mask_hidden_dim,
+            output_dim=hidden_dim,
+        )
         self.ffn = nn.Linear(hidden_dim, hidden_dim)
 
     # different from MaskBlockOnEmbedding: input is hidden representation
-    def forward(self, hidden_in: torch.Tensor, v_emb_flat: torch.Tensor) -> torch.Tensor:
-        norm_hidden = self.ln_input(hidden_in)  
+    def forward(
+        self, hidden_in: torch.Tensor, v_emb_flat: torch.Tensor
+    ) -> torch.Tensor:
+        norm_hidden = self.ln_input(hidden_in)
         v_mask = self.mask_gen(v_emb_flat)
-        v_masked_hid = v_mask * norm_hidden     
-        out = self.ffn(v_masked_hid)  
+        v_masked_hid = v_mask * norm_hidden
+        out = self.ffn(v_masked_hid)
         out = self.ln_output(out)
         out = F.relu(out)
         return out
@@ -151,7 +166,7 @@ class MaskNet(BaseModel):
         dense_features: list[DenseFeature] | None = None,
         sparse_features: list[SparseFeature] | None = None,
         sequence_features: list[SequenceFeature] | None = None,
-        model_type: str = "parallel",         # "serial" or "parallel"
+        architecture: str = "parallel",  # "serial" or "parallel"
         num_blocks: int = 3,
         mask_hidden_dim: int = 64,
         block_hidden_dim: int = 256,
@@ -199,52 +214,99 @@ class MaskNet(BaseModel):
         self.sparse_features = sparse_features
         self.sequence_features = sequence_features
         self.mask_features = self.all_features  # use all features for masking
-        assert len(self.mask_features) > 0, "MaskNet requires at least one feature for masking."
+        assert (
+            len(self.mask_features) > 0
+        ), "MaskNet requires at least one feature for masking."
         self.embedding = EmbeddingLayer(features=self.mask_features)
         self.num_fields = len(self.mask_features)
         self.embedding_dim = getattr(self.mask_features[0], "embedding_dim", None)
-        assert self.embedding_dim is not None, "MaskNet requires mask_features to have 'embedding_dim' defined."
+        assert (
+            self.embedding_dim is not None
+        ), "MaskNet requires mask_features to have 'embedding_dim' defined."
 
         for f in self.mask_features:
             edim = getattr(f, "embedding_dim", None)
             if edim is None or edim != self.embedding_dim:
-                raise ValueError(f"MaskNet expects identical embedding_dim across all mask_features, but got {edim} for feature {getattr(f, 'name', type(f))}.")
+                raise ValueError(
+                    f"MaskNet expects identical embedding_dim across all mask_features, but got {edim} for feature {getattr(f, 'name', type(f))}."
+                )
 
         self.v_emb_dim = self.num_fields * self.embedding_dim
-        self.model_type = model_type.lower()
-        assert self.model_type in ("serial", "parallel"), "model_type must be either 'serial' or 'parallel'."
+        self.architecture = architecture.lower()
+        assert self.architecture in (
+            "serial",
+            "parallel",
+        ), "architecture must be either 'serial' or 'parallel'."
 
         self.num_blocks = max(1, num_blocks)
         self.block_hidden_dim = block_hidden_dim
-        self.block_dropout = nn.Dropout(block_dropout) if block_dropout > 0 else nn.Identity()
+        self.block_dropout = (
+            nn.Dropout(block_dropout) if block_dropout > 0 else nn.Identity()
+        )
 
-        if self.model_type == "serial":
-            self.first_block = MaskBlockOnEmbedding(num_fields=self.num_fields, embedding_dim=self.embedding_dim, mask_hidden_dim=mask_hidden_dim, hidden_dim=block_hidden_dim,)
+        if self.architecture == "serial":
+            self.first_block = MaskBlockOnEmbedding(
+                num_fields=self.num_fields,
+                embedding_dim=self.embedding_dim,
+                mask_hidden_dim=mask_hidden_dim,
+                hidden_dim=block_hidden_dim,
+            )
             self.hidden_blocks = nn.ModuleList(
-                [MaskBlockOnHidden(num_fields=self.num_fields, embedding_dim=self.embedding_dim, mask_hidden_dim=mask_hidden_dim, hidden_dim=block_hidden_dim) for _ in range(self.num_blocks - 1)])
+                [
+                    MaskBlockOnHidden(
+                        num_fields=self.num_fields,
+                        embedding_dim=self.embedding_dim,
+                        mask_hidden_dim=mask_hidden_dim,
+                        hidden_dim=block_hidden_dim,
+                    )
+                    for _ in range(self.num_blocks - 1)
+                ]
+            )
             self.mask_blocks = nn.ModuleList([self.first_block, *self.hidden_blocks])
             self.output_layer = nn.Linear(block_hidden_dim, 1)
             self.final_mlp = None
 
         else:  # parallel
-            self.mask_blocks = nn.ModuleList([MaskBlockOnEmbedding(num_fields=self.num_fields, embedding_dim=self.embedding_dim, mask_hidden_dim=mask_hidden_dim, hidden_dim=block_hidden_dim) for _ in range(self.num_blocks)])
-            self.final_mlp = MLP(input_dim=self.num_blocks * block_hidden_dim, **mlp_params)
+            self.mask_blocks = nn.ModuleList(
+                [
+                    MaskBlockOnEmbedding(
+                        num_fields=self.num_fields,
+                        embedding_dim=self.embedding_dim,
+                        mask_hidden_dim=mask_hidden_dim,
+                        hidden_dim=block_hidden_dim,
+                    )
+                    for _ in range(self.num_blocks)
+                ]
+            )
+            self.final_mlp = MLP(
+                input_dim=self.num_blocks * block_hidden_dim, **mlp_params
+            )
             self.output_layer = None
         self.prediction_layer = PredictionLayer(task_type=self.task)
 
-        if self.model_type == "serial":
-            self.register_regularization_weights(embedding_attr="embedding", include_modules=["mask_blocks", "output_layer"],)
+        if self.architecture == "serial":
+            self.register_regularization_weights(
+                embedding_attr="embedding",
+                include_modules=["mask_blocks", "output_layer"],
+            )
         # serial
         else:
-            self.register_regularization_weights(embedding_attr="embedding", include_modules=["mask_blocks", "final_mlp"])
-        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)
+            self.register_regularization_weights(
+                embedding_attr="embedding", include_modules=["mask_blocks", "final_mlp"]
+            )
+        self.compile(
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            loss=loss,
+            loss_params=loss_params,
+        )
 
     def forward(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
         field_emb = self.embedding(x=x, features=self.mask_features, squeeze_dim=False)
         B = field_emb.size(0)
-        v_emb_flat = field_emb.view(B, -1)  # flattened embeddings 
+        v_emb_flat = field_emb.view(B, -1)  # flattened embeddings
 
-        if self.model_type == "parallel":
+        if self.architecture == "parallel":
             block_outputs = []
             for block in self.mask_blocks:
                 h = block(field_emb, v_emb_flat)  # [B, block_hidden_dim]
@@ -253,7 +315,7 @@ class MaskNet(BaseModel):
             concat_hidden = torch.cat(block_outputs, dim=-1)
             logit = self.final_mlp(concat_hidden)  # [B, 1]
         # serial
-        else: 
+        else:
             hidden = self.first_block(field_emb, v_emb_flat)
             hidden = self.block_dropout(hidden)
             for block in self.hidden_blocks:

nextrec/models/ranking/pnn.py

@@ -3,7 +3,35 @@ Date: create on 09/11/2025
 Author:
     Yang Zhou,zyaztec@gmail.com
 Reference:
-    [1] Qu Y, Cai H, Ren K, et al. Product-based neural networks for user response prediction[C]//ICDM. 2016: 1149-1154.
+[1] Qu Y, Cai H, Ren K, et al. Product-based neural networks for user response
+prediction[C]//ICDM. 2016: 1149-1154. (https://arxiv.org/abs/1611.00144)
+
+Product-based Neural Networks (PNN) are CTR prediction models that explicitly
+encode feature interactions by combining:
+  (1) A linear signal from concatenated field embeddings
+  (2) A product signal capturing pairwise feature interactions (inner or outer)
+The product layer augments the linear input to an MLP, enabling the network to
+model both first-order and high-order feature interactions in a structured way.
+
+Computation workflow:
+  - Embed each categorical/sequence field with a shared embedding dimension
+  - Linear signal: flatten and concatenate all field embeddings
+  - Product signal:
+      * Inner product: dot products over all field pairs
+      * Outer product: project embeddings then compute element-wise products
+  - Concatenate linear and product signals; feed into MLP for prediction
+
+Key Advantages:
+- Explicit pairwise interaction modeling without heavy feature engineering
+- Flexible choice between inner/outer products to trade off capacity vs. cost
+- Combines linear context with interaction signal for stronger expressiveness
+- Simple architecture that integrates cleanly with standard MLP pipelines
+
+PNN 是一种 CTR 预估模型，通过将线性信号与乘积信号结合，显式建模特征交互：
+  - 线性信号：将各字段的 embedding 拼接，用于保留一阶信息
+  - 乘积信号：对所有字段对做内积或外积，捕捉二阶及更高阶交互
+随后将两类信号拼接送入 MLP，实现对用户响应的预测。内积版本计算量更低，
+外积版本表达力更强，可根据场景取舍。
 """
 
 import torch
@@ -15,6 +43,7 @@ from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 
 
 class PNN(BaseModel):
+
     @property
     def model_name(self):
         return "PNN"
@@ -22,27 +51,39 @@ class PNN(BaseModel):
     @property
     def default_task(self):
         return "binary"
-    
-    def __init__(self,
-                 dense_features: list[DenseFeature] | list = [],
-                 sparse_features: list[SparseFeature] | list = [],
-                 sequence_features: list[SequenceFeature] | list = [],
-                 mlp_params: dict = {},
-                 product_type: str = "inner",
-                 outer_product_dim: int | None = None,
-                 target: list[str] | list = [],
-                 task: str | list[str] | None = None,
-                 optimizer: str = "adam",
-                 optimizer_params: dict = {},
-                 loss: str | nn.Module | None = "bce",
-                 loss_params: dict | list[dict] | None = None,
-                 device: str = 'cpu',
-                 embedding_l1_reg=1e-6,
-                 dense_l1_reg=1e-5,
-                 embedding_l2_reg=1e-5,
-                 dense_l2_reg=1e-4,
-                 **kwargs):
-        
+
+    def __init__(
+        self,
+        dense_features: list[DenseFeature] | None = None,
+        sparse_features: list[SparseFeature] | None = None,
+        sequence_features: list[SequenceFeature] | None = None,
+        mlp_params: dict | None = None,
+        product_type: str = "inner",  # "inner" (IPNN), "outer" (OPNN), "both" (PNN*)
+        outer_product_dim: int | None = None,
+        target: list[str] | str | None = None,
+        task: str | list[str] | None = None,
+        optimizer: str = "adam",
+        optimizer_params: dict | None = None,
+        loss: str | nn.Module | None = "bce",
+        loss_params: dict | list[dict] | None = None,
+        device: str = "cpu",
+        embedding_l1_reg=1e-6,
+        dense_l1_reg=1e-5,
+        embedding_l2_reg=1e-5,
+        dense_l2_reg=1e-4,
+        **kwargs,
+    ):
+
+        dense_features = dense_features or []
+        sparse_features = sparse_features or []
+        sequence_features = sequence_features or []
+        mlp_params = mlp_params or {}
+        if outer_product_dim is not None and outer_product_dim <= 0:
+            raise ValueError("outer_product_dim must be a positive integer.")
+        optimizer_params = optimizer_params or {}
+        if loss is None:
+            loss = "bce"
+
         super(PNN, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -54,46 +95,54 @@ class PNN(BaseModel):
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            **kwargs
+            **kwargs,
         )
 
-        self.loss = loss
-        if self.loss is None:
-            self.loss = "bce"
-            
-        self.field_features = sparse_features + sequence_features
+        self.field_features = dense_features + sparse_features + sequence_features
         if len(self.field_features) < 2:
             raise ValueError("PNN requires at least two sparse/sequence features.")
 
         self.embedding = EmbeddingLayer(features=self.field_features)
         self.num_fields = len(self.field_features)
+
         self.embedding_dim = self.field_features[0].embedding_dim
         if any(f.embedding_dim != self.embedding_dim for f in self.field_features):
-            raise ValueError("All field features must share the same embedding_dim for PNN.")
+            raise ValueError(
+                "All field features must share the same embedding_dim for PNN."
+            )
 
         self.product_type = product_type.lower()
-        if self.product_type not in {"inner", "outer"}:
-            raise ValueError("product_type must be 'inner' or 'outer'.")
+        if self.product_type not in {"inner", "outer", "both"}:
+            raise ValueError("product_type must be 'inner', 'outer', or 'both'.")
 
         self.num_pairs = self.num_fields * (self.num_fields - 1) // 2
-        if self.product_type == "outer":
-            self.outer_dim = outer_product_dim or self.embedding_dim
-            self.kernel = nn.Linear(self.embedding_dim, self.outer_dim, bias=False)
-            product_dim = self.num_pairs * self.outer_dim
+        self.outer_product_dim = outer_product_dim or self.embedding_dim
+
+        if self.product_type in {"outer", "both"}:
+            self.kernel = nn.Parameter(
+                torch.randn(self.embedding_dim, self.outer_product_dim)
+            )
+            nn.init.xavier_uniform_(self.kernel)
         else:
-            self.outer_dim = None
-            product_dim = self.num_pairs
+            self.kernel = None
 
         linear_dim = self.num_fields * self.embedding_dim
+
+        if self.product_type == "inner":
+            product_dim = self.num_pairs
+        elif self.product_type == "outer":
+            product_dim = self.num_pairs
+        else:
+            product_dim = 2 * self.num_pairs
+
         self.mlp = MLP(input_dim=linear_dim + product_dim, **mlp_params)
         self.prediction_layer = PredictionLayer(task_type=self.task)
 
-        modules = ['mlp']
-        if self.product_type == "outer":
-            modules.append('kernel')
+        modules = ["mlp"]
+        if self.kernel is not None:
+            modules.append("kernel")
         self.register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=modules
+            embedding_attr="embedding", include_modules=modules
         )
 
         self.compile(
@@ -103,27 +152,48 @@ class PNN(BaseModel):
             loss_params=loss_params,
         )
 
+    def compute_inner_products(self, field_emb: torch.Tensor) -> torch.Tensor:
+        interactions = []
+        for i in range(self.num_fields - 1):
+            vi = field_emb[:, i, :]  # [B, D]
+            for j in range(i + 1, self.num_fields):
+                vj = field_emb[:, j, :]  # [B, D]
+                # <v_i, v_j> = sum_k v_i,k * v_j,k
+                pij = torch.sum(vi * vj, dim=1, keepdim=True)  # [B, 1]
+                interactions.append(pij)
+        return torch.cat(interactions, dim=1)  # [B, num_pairs]
+
+    def compute_outer_kernel_products(self, field_emb: torch.Tensor) -> torch.Tensor:
+        if self.kernel is None:
+            raise RuntimeError("kernel is not initialized for outer product.")
+
+        interactions = []
+        for i in range(self.num_fields - 1):
+            vi = field_emb[:, i, :]  # [B, D]
+            # Project vi with kernel -> [B, K]
+            vi_proj = torch.matmul(vi, self.kernel)  # [B, K]
+            for j in range(i + 1, self.num_fields):
+                vj = field_emb[:, j, :]  # [B, D]
+                vj_proj = torch.matmul(vj, self.kernel)  # [B, K]
+                # g(vi, vj) = (v_i^T W) * (v_j^T W) summed over projection dim
+                pij = torch.sum(vi_proj * vj_proj, dim=1, keepdim=True)  # [B, 1]
+                interactions.append(pij)
+        return torch.cat(interactions, dim=1)  # [B, num_pairs]
+
     def forward(self, x):
+        # field_emb: [B, F, D]
         field_emb = self.embedding(x=x, features=self.field_features, squeeze_dim=False)
-        linear_signal = field_emb.flatten(start_dim=1)
+        # Z = [v_1; v_2; ...; v_F]
+        linear_signal = field_emb.flatten(start_dim=1)  # [B, F*D]
 
         if self.product_type == "inner":
-            interactions = []
-            for i in range(self.num_fields - 1):
-                vi = field_emb[:, i, :]
-                for j in range(i + 1, self.num_fields):
-                    vj = field_emb[:, j, :]
-                    interactions.append(torch.sum(vi * vj, dim=1, keepdim=True))
-            product_signal = torch.cat(interactions, dim=1)
+            product_signal = self.compute_inner_products(field_emb)
+        elif self.product_type == "outer":
+            product_signal = self.compute_outer_kernel_products(field_emb)
         else:
-            transformed = self.kernel(field_emb)  # [B, F, outer_dim]
-            interactions = []
-            for i in range(self.num_fields - 1):
-                vi = transformed[:, i, :]
-                for j in range(i + 1, self.num_fields):
-                    vj = transformed[:, j, :]
-                    interactions.append(vi * vj)
-            product_signal = torch.stack(interactions, dim=1).flatten(start_dim=1)
+            inner_p = self.compute_inner_products(field_emb)
+            outer_p = self.compute_outer_kernel_products(field_emb)
+            product_signal = torch.cat([inner_p, outer_p], dim=1)
 
         deep_input = torch.cat([linear_signal, product_signal], dim=1)
         y = self.mlp(deep_input)