nextrec 0.2.7__py3-none-any.whl → 0.3.1__py3-none-any.whl

nextrec/models/multi_task/poso.py

@@ -0,0 +1,413 @@
+"""
+Date: create on 28/11/2025
+Author: Yang Zhou,zyaztec@gmail.com
+Reference:
+[1] Wang et al. "POSO: Personalized Cold Start Modules for Large-scale Recommender Systems", 2021.
+
+POSO (Personalized cOld-start mOdules) augments backbone recommenders by injecting a
+personalized cold-start vector `pc` that gates hidden units layer by layer. Each fully
+connected layer or expert output is multiplied by gate(pc), letting the backbone adapt
+its hidden representations to user profiles even when behavioral signals are scarce.
+
+Core idea:
+  (1) A lightweight two-layer MLP maps `pc` to gate(pc) = C * sigmoid(W2 * phi(W1 * pc + b1) + b2)
+  (2) gate(pc) scales each hidden unit element-wise, masking or amplifying features
+  (3) Existing task gates/towers remain intact; POSO only overlays personalization
+
+Key advantages:
+- Plug-and-play personalization for cold-start users without redesigning the backbone
+- Per-layer/expert gating with minimal additional parameters
+- Compatible with plain MLP towers and MMoE structures, keeping training stable
+- Works with split features: main features feed the backbone, PC features drive gates
+
+POSO 通过个性化冷启动向量 `pc` 为推荐模型叠加逐层的门控系数，
+在每个全连接层或专家输出上乘以 gate(pc) 做元素级缩放，
+即使行为信号稀缺也能按用户画像调整隐藏表示。
+
+实现思路：
+  (1) 用轻量两层 MLP 生成 gate(pc) = C * sigmoid(W2 * phi(W1 * pc + b1) + b2)
+  (2) gate(pc) 对神经元逐元素放大或抑制
+  (3) 原有任务门/塔不变，POSO 仅叠加个性化门控
+
+主要优点：
+- 冷启动场景的可插拔个性化，无需重做骨干结构
+- 每层/每专家独立门控，新增参数量小
+- 兼容 MLP、MMoE 等多任务骨干，训练过程平稳
+- 主特征做建模，PC 特征驱动门控，解耦表征与个性化信号
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
+from nextrec.basic.layers import EmbeddingLayer, MLP, PredictionLayer
+from nextrec.basic.activation import activation_layer
+from nextrec.basic.model import BaseModel
+from nextrec.utils.common import merge_features
+
+
+class POSOGate(nn.Module):
+    """
+    Two-layer MLP that maps the personalized cold-start vector to a gate:
+        gate(pc) = C * sigmoid( W2 * phi(W1 * pc + b1) + b2 )
+    The output shares the same dimension as the hidden vector to be masked and
+    is applied element-wise.
+    """
+
+    def __init__(
+        self,
+        pc_dim: int,
+        out_dim: int,
+        hidden_dim: int = 32,
+        scale_factor: float = 2.0,
+        activation: str = "relu",
+    ) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(pc_dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, out_dim)
+        self.scale_factor = scale_factor
+        self.act = activation_layer(activation)
+
+    def forward(self, pc: torch.Tensor) -> torch.Tensor:
+        """
+        pc: (B, pc_dim)
+        return: (B, out_dim) in (0, C)
+        """
+        h = self.act(self.fc1(pc))
+        g = torch.sigmoid(self.fc2(h))  # (B, out_dim) in (0,1)
+        return self.scale_factor * g
+    
+
+class POSOFC(nn.Module):
+    """
+    Single POSO fully connected layer mirroring Eq. (11):
+        h = phi(Wx + b)
+        h_hat = gate(pc) ⊙ h
+    where gate(pc) = C * sigmoid(MLP(pc)).
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        pc_dim: int,
+        gate_hidden_dim: int = 32,
+        scale_factor: float = 2.0,
+        activation: str = "relu",
+        use_bias: bool = True,
+    ) -> None:
+        super().__init__()
+        self.linear = nn.Linear(in_dim, out_dim, bias=use_bias)
+        self.act = activation_layer(activation)
+        self.gate = POSOGate(
+            pc_dim=pc_dim,
+            out_dim=out_dim,
+            hidden_dim=gate_hidden_dim,
+            scale_factor=scale_factor,
+            activation=activation,
+        )
+
+    def forward(self, x: torch.Tensor, pc: torch.Tensor) -> torch.Tensor:
+        """
+        x:  (B, in_dim)
+        pc: (B, pc_dim)
+        return: (B, out_dim)
+        """
+        h = self.act(self.linear(x))          # Standard FC with activation
+        g = self.gate(pc)                     # (B, out_dim)
+        return g * h                          # Element-wise gating
+    
+
+class POSOMLP(nn.Module):
+    """
+    POSO-enhanced MLP that stacks multiple POSOFC layers.
+
+    dims: e.g., [256, 128, 64] means
+        in_dim -> 256 -> 128 -> 64
+    Each layer has its own gate g_l(pc) following Eq. (11).
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        pc_dim: int,
+        dims: list[int],
+        gate_hidden_dim: int = 32,
+        scale_factor: float = 2.0,
+        activation: str = "relu",
+        use_bias: bool = True,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+
+        layers = []
+        in_dim = input_dim
+        for out_dim in dims:
+            layers.append(
+                POSOFC(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    pc_dim=pc_dim,
+                    gate_hidden_dim=gate_hidden_dim,
+                    scale_factor=scale_factor,
+                    activation=activation,
+                    use_bias=use_bias,
+                )
+            )
+            in_dim = out_dim
+
+        self.layers = nn.ModuleList(layers)
+        self.dropout = nn.Dropout(dropout) if dropout > 0 else None
+
+    def forward(self, x: torch.Tensor, pc: torch.Tensor) -> torch.Tensor:
+        """
+        x:  (B, input_dim)
+        pc: (B, pc_dim)
+        """
+        h = x
+        for layer in self.layers:
+            h = layer(h, pc)
+            if self.dropout is not None:
+                h = self.dropout(h)
+        return h
+    
+
+class POSOMMoE(nn.Module):
+    """
+    POSO(MMoE) mirrors Section 4.4 and Eq. (15)-(18) of the paper:
+        - Keep the original experts and task gates gate_t(x)
+        - Add a PC gate g_e(pc) for every expert_e
+        - Task gates aggregate the PC-masked expert outputs
+
+    Concretely:
+        h_e = expert_e(x)                 # (B, D) 
+        g_e = POSOGate(pc) in (0, C)^{D}  # (B, D)
+        h_e_tilde = g_e ⊙ h_e            # (B, D)
+        z_t = Σ_e gate_t,e(x) * h_e_tilde
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        pc_dim: int,                   # for poso feature dimension
+        num_experts: int,
+        expert_hidden_dims: list[int],
+        num_tasks: int,
+        activation: str = "relu",
+        expert_dropout: float = 0.0,
+        gate_hidden_dim: int = 32,     # for poso gate hidden dimension
+        scale_factor: float = 2.0,     # for poso gate scale factor
+        gate_use_softmax: bool = True,
+    ) -> None:
+        super().__init__()
+        self.num_experts = num_experts
+        self.num_tasks = num_tasks
+
+        # Experts built with framework MLP, same as standard MMoE
+        self.experts = nn.ModuleList([MLP(input_dim=input_dim, output_layer=False, dims=expert_hidden_dims, activation=activation, dropout=expert_dropout,) for _ in range(num_experts)])
+        self.expert_output_dim = expert_hidden_dims[-1] if expert_hidden_dims else input_dim
+
+        # Task-specific gates: gate_t(x) over experts
+        self.gates = nn.ModuleList([nn.Linear(input_dim, num_experts) for _ in range(num_tasks)])
+        self.gate_use_softmax = gate_use_softmax
+
+        # PC gate per expert: g_e(pc) ∈ R^D
+        self.expert_pc_gates = nn.ModuleList([POSOGate(pc_dim=pc_dim, out_dim=self.expert_output_dim, hidden_dim=gate_hidden_dim, scale_factor=scale_factor, activation=activation,) for _ in range(num_experts)])
+
+    def forward(self, x: torch.Tensor, pc: torch.Tensor) -> list[torch.Tensor]:
+        """
+        x:  (B, input_dim)
+        pc: (B, pc_dim)
+        return: list of task outputs z_t with length num_tasks, each (B, D)
+        """
+        # 1) Expert outputs with POSO PC gate
+        masked_expert_outputs = []
+        for e, expert in enumerate(self.experts):
+            h_e = expert(x)                         # (B, D)
+            g_e = self.expert_pc_gates[e](pc)       # (B, D)
+            h_e_tilde = g_e * h_e                   # (B, D)
+            masked_expert_outputs.append(h_e_tilde)
+
+        masked_expert_outputs = torch.stack(masked_expert_outputs, dim=1)  # (B, E, D)
+
+        # 2) Task gates depend on x as in standard MMoE
+        task_outputs: list[torch.Tensor] = []
+        for t in range(self.num_tasks):
+            logits = self.gates[t](x)               # (B, E)
+            if self.gate_use_softmax:
+                gate = F.softmax(logits, dim=1)
+            else:
+                gate = logits
+
+            gate = gate.unsqueeze(-1)               # (B, E, 1)
+            z_t = torch.sum(gate * masked_expert_outputs, dim=1)  # (B, D)
+            task_outputs.append(z_t)
+
+        return task_outputs
+
+
+class POSO(BaseModel):
+    """
+    POSO model implemented with the NextRec framework. It supports two backbones:
+    - "mlp": per-task POSO-MLP towers with PC gating on every hidden layer
+    - "mmoe": POSO-gated MMoE experts plus task-specific towers
+    """
+
+    @property
+    def model_name(self) -> str:
+        return "POSO"
+
+    @property
+    def task_type(self) -> list[str]:
+        return self.task if isinstance(self.task, list) else [self.task]
+
+    def __init__(
+        self,
+        main_dense_features: list[DenseFeature] | None,
+        main_sparse_features: list[SparseFeature] | None,
+        main_sequence_features: list[SequenceFeature] | None,
+        pc_dense_features: list[DenseFeature] | None,
+        pc_sparse_features: list[SparseFeature] | None,
+        pc_sequence_features: list[SequenceFeature] | None,
+        tower_params_list: list[dict],
+        target: list[str],
+        task: str | list[str] = "binary",
+        architecture: str = "mlp",
+        # POSO gating defaults
+        gate_hidden_dim: int = 32,
+        gate_scale_factor: float = 2.0,
+        gate_activation: str = "relu",
+        gate_use_bias: bool = True,
+        # MMoE-specific params
+        num_experts: int = 4,
+        expert_hidden_dims: list[int] | None = None,
+        expert_activation: str = "relu",
+        expert_dropout: float = 0.0,
+        expert_gate_hidden_dim: int = 32,
+        expert_gate_scale_factor: float = 2.0,
+        gate_use_softmax: bool = True,
+        optimizer: str = "adam",
+        optimizer_params: dict | None = None,
+        loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+        loss_params: dict | list[dict] | None = None,
+        device: str = "cpu",
+        embedding_l1_reg: float = 1e-6,
+        dense_l1_reg: float = 1e-5,
+        embedding_l2_reg: float = 1e-5,
+        dense_l2_reg: float = 1e-4,
+        **kwargs,
+    ):
+        # Keep explicit copies of main and PC features
+        self.main_dense_features = list(main_dense_features or [])
+        self.main_sparse_features = list(main_sparse_features or [])
+        self.main_sequence_features = list(main_sequence_features or [])
+        self.pc_dense_features = list(pc_dense_features or [])
+        self.pc_sparse_features = list(pc_sparse_features or [])
+        self.pc_sequence_features = list(pc_sequence_features or [])
+
+        if not self.pc_dense_features and not self.pc_sparse_features and not self.pc_sequence_features:
+            raise ValueError("POSO requires at least one PC feature for personalization.")
+
+        dense_features = merge_features(self.main_dense_features, self.pc_dense_features)
+        sparse_features = merge_features(self.main_sparse_features, self.pc_sparse_features)
+        sequence_features = merge_features(self.main_sequence_features, self.pc_sequence_features)
+
+        super().__init__(
+            dense_features=dense_features,
+            sparse_features=sparse_features,
+            sequence_features=sequence_features,
+            target=target,
+            task=task,
+            device=device,
+            embedding_l1_reg=embedding_l1_reg,
+            dense_l1_reg=dense_l1_reg,
+            embedding_l2_reg=embedding_l2_reg,
+            dense_l2_reg=dense_l2_reg,
+            early_stop_patience=20,
+            **kwargs,
+        )
+
+        self.loss = loss if loss is not None else "bce"
+        optimizer_params = optimizer_params or {}
+
+        self.num_tasks = len(target)
+        if len(tower_params_list) != self.num_tasks:
+            raise ValueError(f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.num_tasks})")
+
+        self.main_features = self.main_dense_features + self.main_sparse_features + self.main_sequence_features
+        self.pc_features = self.pc_dense_features + self.pc_sparse_features + self.pc_sequence_features
+
+        self.embedding = EmbeddingLayer(features=self.all_features)
+        self.main_input_dim = self.embedding.get_input_dim(self.main_features)
+        self.pc_input_dim = self.embedding.get_input_dim(self.pc_features)
+
+        self.architecture = architecture.lower()
+        if self.architecture not in {"mlp", "mmoe"}:
+            raise ValueError(f"Unsupported architecture '{architecture}', choose from ['mlp', 'mmoe'].")
+
+        # Build backbones
+        if self.architecture == "mlp":
+            self.towers = nn.ModuleList()
+            self.tower_heads = nn.ModuleList()
+            for tower_params in tower_params_list:
+                dims = tower_params.get("dims")
+                if not dims:
+                    raise ValueError("tower_params must include a non-empty 'dims' list for POSO-MLP towers.")
+                dropout = tower_params.get("dropout", 0.0)
+                tower = POSOMLP(
+                    input_dim=self.main_input_dim,
+                    pc_dim=self.pc_input_dim,
+                    dims=dims,
+                    gate_hidden_dim=tower_params.get("gate_hidden_dim", gate_hidden_dim),
+                    scale_factor=tower_params.get("scale_factor", gate_scale_factor),
+                    activation=tower_params.get("activation", gate_activation),
+                    use_bias=tower_params.get("use_bias", gate_use_bias),
+                    dropout=dropout,
+                )
+                self.towers.append(tower)
+                tower_output_dim = dims[-1] if dims else self.main_input_dim
+                self.tower_heads.append(nn.Linear(tower_output_dim, 1))
+        else:
+            if expert_hidden_dims is None or not expert_hidden_dims:
+                raise ValueError("expert_hidden_dims must be provided for MMoE architecture.")
+            self.mmoe = POSOMMoE(
+                input_dim=self.main_input_dim,
+                pc_dim=self.pc_input_dim,
+                num_experts=num_experts,
+                expert_hidden_dims=expert_hidden_dims,
+                num_tasks=self.num_tasks,
+                activation=expert_activation,
+                expert_dropout=expert_dropout,
+                gate_hidden_dim=expert_gate_hidden_dim,
+                scale_factor=expert_gate_scale_factor,
+                gate_use_softmax=gate_use_softmax,
+            )
+            self.towers = nn.ModuleList([MLP(input_dim=self.mmoe.expert_output_dim, output_layer=True, **tower_params,) for tower_params in tower_params_list])
+            self.tower_heads = None
+        self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1] * self.num_tasks,)
+        include_modules = ["towers", "tower_heads"] if self.architecture == "mlp" else ["mmoe", "towers"]
+        self._register_regularization_weights(embedding_attr="embedding", include_modules=include_modules)
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)
+
+    def forward(self, x):
+        # Embed main and PC features separately so PC can gate hidden units
+        main_input = self.embedding(x=x, features=self.main_features, squeeze_dim=True)
+        pc_input = self.embedding(x=x, features=self.pc_features, squeeze_dim=True)
+
+        task_outputs = []
+        if self.architecture == "mlp":
+            for tower, head in zip(self.towers, self.tower_heads):
+                hidden = tower(main_input, pc_input)
+                logit = head(hidden)
+                task_outputs.append(logit)
+        else:
+            expert_outputs = self.mmoe(main_input, pc_input)
+            for idx, tower in enumerate(self.towers):
+                logit = tower(expert_outputs[idx])
+                task_outputs.append(logit)
+
+        y = torch.cat(task_outputs, dim=1)
+        return self.prediction_layer(y)

nextrec/models/multi_task/share_bottom.py

@@ -1,7 +1,42 @@
 """
 Date: create on 09/11/2025
+Checkpoint: edit on 24/11/2025
 Author: Yang Zhou,zyaztec@gmail.com
-Reference: [1] Caruana R. Multitask learning[J]. Machine learning, 1997, 28: 41-75.
+Reference:
+[1] Caruana R. Multitask learning[J]. Machine Learning, 1997, 28: 41-75.
+(https://link.springer.com/article/10.1023/A:1007379606734)
+
+Shared-Bottom is the classic hard-parameter-sharing baseline for multi-task learning.
+All tasks share a common bottom network to learn general representations, and each
+task has its own tower head for task-specific refinement and prediction. This
+architecture is simple, parameter-efficient, and helps regularize related tasks.
+
+Workflow:
+  (1) Unified embeddings convert dense/sparse/sequence features
+  (2) A shared bottom MLP learns common representations
+  (3) Task-specific towers further transform the shared features
+  (4) Separate prediction heads output each task’s logits/probabilities
+
+Key Advantages:
+- Strong inductive bias via hard parameter sharing, reducing overfitting
+- Parameter-efficient compared to duplicating full models per task
+- Easy to extend to many tasks with small incremental cost
+- Serves as a stable baseline for evaluating advanced MTL architectures
+
+Share-Bottom（硬共享底层）是多任务学习的经典基线：所有任务共享一个底层网络，
+各任务拥有独立塔头进行细化与预测，简单高效且能通过共享正则化相关任务。
+
+流程：
+  (1) 统一 embedding 处理稠密、稀疏与序列特征
+  (2) 共享底层 MLP 学习通用表示
+  (3) 任务塔在共享表示上做任务特定变换
+  (4) 各任务预测头输出对应结果
+
+主要优点：
+- 硬参数共享提供强正则，减少过拟合
+- 相比单独模型更节省参数与计算
+- 易于扩展到多任务，增量开销小
+- 是评估更复杂 MTL 结构的稳健基线
 """
 
 import torch
@@ -59,22 +94,17 @@ class ShareBottom(BaseModel):
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-        
         # Number of tasks
         self.num_tasks = len(target)
         if len(tower_params_list) != self.num_tasks:
             raise ValueError(f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.num_tasks})")
-            
-        # All features
-        self.all_features = dense_features + sparse_features + sequence_features
-
         # Embedding layer
         self.embedding = EmbeddingLayer(features=self.all_features)
-
         # Calculate input dimension
-        emb_dim_total = sum([f.embedding_dim for f in self.all_features if not isinstance(f, DenseFeature)])
-        dense_input_dim = sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
-        input_dim = emb_dim_total + dense_input_dim
+        input_dim = self.embedding.input_dim
+        # emb_dim_total = sum([f.embedding_dim for f in self.all_features if not isinstance(f, DenseFeature)])
+        # dense_input_dim = sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
+        # input_dim = emb_dim_total + dense_input_dim
         
         # Shared bottom network
         self.bottom = MLP(input_dim=input_dim, output_layer=False, **bottom_params)
@@ -90,23 +120,10 @@ class ShareBottom(BaseModel):
         for tower_params in tower_params_list:
             tower = MLP(input_dim=bottom_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
-        self.prediction_layer = PredictionLayer(
-            task_type=self.task_type,
-            task_dims=[1] * self.num_tasks
-        )
-
+        self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1] * self.num_tasks)
         # Register regularization weights
-        self._register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=['bottom', 'towers']
-        )
-
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss,
-            loss_params=loss_params,
-        )
+        self._register_regularization_weights(embedding_attr='embedding', include_modules=['bottom', 'towers'])
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)
 
     def forward(self, x):
         # Get all embeddings and flatten

nextrec/models/ranking/__init__.py

@@ -7,6 +7,7 @@ from .autoint import AutoInt
 from .widedeep import WideDeep
 from .xdeepfm import xDeepFM
 from .dcn import DCN
+from .fibinet import FiBiNET
 from .din import DIN
 from .dien import DIEN
 
@@ -22,4 +23,5 @@ __all__ = [
     'AFM',
     'MaskNet',
     'PNN',
+    'FiBiNET',
 ]

nextrec/models/ranking/dcn.py

@@ -12,9 +12,28 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.model import BaseModel
-from nextrec.basic.layers import EmbeddingLayer, MLP, CrossNetwork, PredictionLayer
+from nextrec.basic.layers import EmbeddingLayer, MLP, PredictionLayer
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 
+class CrossNetwork(nn.Module):
+    """Stacked Cross Layers from DCN (Wang et al., 2017)."""
+
+    def __init__(self, input_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        self.w = torch.nn.ModuleList([torch.nn.Linear(input_dim, 1, bias=False) for _ in range(num_layers)])
+        self.b = torch.nn.ParameterList([torch.nn.Parameter(torch.zeros((input_dim,))) for _ in range(num_layers)])
+
+    def forward(self, x):
+        """
+        :param x: Float tensor of size ``(batch_size, num_fields, embed_dim)``
+        """
+        x0 = x
+        for i in range(self.num_layers):
+            xw = self.w[i](x)
+            x = x0 * xw + self.b[i] + x
+        return x
+
 
 class DCN(BaseModel):
     @property

nextrec/models/ranking/dcn_v2.py

@@ -0,0 +1,84 @@
+"""
+Date: create on 09/11/2025
+"""
+
+import torch
+import torch.nn as nn
+
+from nextrec.basic.model import BaseModel
+from nextrec.basic.layers import EmbeddingLayer, MLP, PredictionLayer
+from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
+
+class CrossNetV2(nn.Module):
+    """Vector-wise cross network proposed in DCN V2 (Wang et al., 2021)."""
+    def __init__(self, input_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        self.w = torch.nn.ModuleList([torch.nn.Linear(input_dim, input_dim, bias=False) for _ in range(num_layers)])
+        self.b = torch.nn.ParameterList([torch.nn.Parameter(torch.zeros((input_dim,))) for _ in range(num_layers)])
+
+
+    def forward(self, x):
+        x0 = x
+        for i in range(self.num_layers):
+            x =x0*self.w[i](x) + self.b[i] + x
+        return x
+    
+class CrossNetMix(nn.Module):
+    """Mixture of low-rank cross experts from DCN V2 (Wang et al., 2021)."""
+
+    def __init__(self, input_dim, num_layers=2, low_rank=32, num_experts=4):
+        super(CrossNetMix, self).__init__()
+        self.num_layers = num_layers
+        self.num_experts = num_experts
+
+        # U: (input_dim, low_rank)
+        self.u_list = torch.nn.ParameterList([nn.Parameter(nn.init.xavier_normal_(
+            torch.empty(num_experts, input_dim, low_rank))) for i in range(self.num_layers)])
+        # V: (input_dim, low_rank)
+        self.v_list = torch.nn.ParameterList([nn.Parameter(nn.init.xavier_normal_(
+            torch.empty(num_experts, input_dim, low_rank))) for i in range(self.num_layers)])
+        # C: (low_rank, low_rank)
+        self.c_list = torch.nn.ParameterList([nn.Parameter(nn.init.xavier_normal_(
+            torch.empty(num_experts, low_rank, low_rank))) for i in range(self.num_layers)])
+        self.gating = nn.ModuleList([nn.Linear(input_dim, 1, bias=False) for i in range(self.num_experts)])
+
+        self.bias = torch.nn.ParameterList([nn.Parameter(nn.init.zeros_(
+            torch.empty(input_dim, 1))) for i in range(self.num_layers)])
+
+    def forward(self, x):
+        x_0 = x.unsqueeze(2)  # (bs, in_features, 1)
+        x_l = x_0
+        for i in range(self.num_layers):
+            output_of_experts = []
+            gating_score_experts = []
+            for expert_id in range(self.num_experts):
+                # (1) G(x_l)
+                # compute the gating score by x_l
+                gating_score_experts.append(self.gating[expert_id](x_l.squeeze(2)))
+
+                # (2) E(x_l)
+                # project the input x_l to $\mathbb{R}^{r}$
+                v_x = torch.matmul(self.v_list[i][expert_id].t(), x_l)  # (bs, low_rank, 1)
+
+                # nonlinear activation in low rank space
+                v_x = torch.tanh(v_x)
+                v_x = torch.matmul(self.c_list[i][expert_id], v_x)
+                v_x = torch.tanh(v_x)
+
+                # project back to $\mathbb{R}^{d}$
+                uv_x = torch.matmul(self.u_list[i][expert_id], v_x)  # (bs, in_features, 1)
+
+                dot_ = uv_x + self.bias[i]
+                dot_ = x_0 * dot_  # Hadamard-product
+
+                output_of_experts.append(dot_.squeeze(2))
+
+            # (3) mixture of low-rank experts
+            output_of_experts = torch.stack(output_of_experts, 2)  # (bs, in_features, num_experts)
+            gating_score_experts = torch.stack(gating_score_experts, 1)  # (bs, num_experts, 1)
+            moe_out = torch.matmul(output_of_experts, gating_score_experts.softmax(1))
+            x_l = moe_out + x_l  # (bs, in_features, 1)
+
+        x_l = x_l.squeeze()  # (bs, in_features)
+        return x_l