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

nextrec/models/generative/hstu.py

@@ -0,0 +1,399 @@
+"""
+[Info: this version is not released yet, i need to more research on source code and paper]
+Date: create on 01/12/2025
+Checkpoint: edit on 01/12/2025 
+Author: Yang Zhou, zyaztec@gmail.com
+Reference:
+[1] Meta AI. Generative Recommenders (HSTU encoder) — https://github.com/meta-recsys/generative-recommenders
+[2] Ma W, Li P, Chen C, et al. Actions speak louder than words: Trillion-parameter sequential transducers for generative recommendations. arXiv:2402.17152.
+
+Hierarchical Sequential Transduction Unit (HSTU) is the core encoder behind 
+Meta’s Generative Recommenders. It replaces softmax attention with lightweight 
+pointwise activations, enabling extremely deep stacks on long behavior sequences.
+
+In each HSTU layer:
+  (1) Tokens are projected into four streams U, V, Q, K via a shared feed-forward block
+  (2) Softmax-free interactions combine QK^T with Relative Attention Bias (RAB) to encode distance
+  (3) Aggregated context is modulated by U-gating and mapped back through an output projection
+
+Stacking layers yields an efficient causal encoder for next-item 
+generation. With a tied-embedding LM head, HSTU forms 
+a full generative recommendation model.
+
+Key Advantages:
+- Softmax-free attention scales better on deep/long sequences
+- RAB captures temporal structure without extra attention heads
+- Causal masking and padding-aware normalization fit real logs
+- Weight tying reduces parameters and stabilizes training
+- Serves as a drop-in backbone for generative recommendation
+
+HSTU（层次化序列转导单元）是 Meta 生成式推荐的核心编码器，
+用点式激活替代 softmax 注意力，可在长序列上轻松堆叠深层结构。
+
+单层 HSTU 的主要步骤：
+  (1) 将输入一次性映射到 U、V、Q、K 四条通路
+  (2) 利用不含 softmax 的 QK^T 结合相对位置偏置（RAB）建模距离信息
+  (3) 用 U 对聚合上下文进行门控，再映射回输出空间
+
+多层堆叠后，可得到高效的因果编码器；与绑权 LM 头配合即可完成 next-item 预测。
+
+主要优势：
+- 摆脱 softmax，在长序列、深层模型上更易扩展
+- 相对位置偏置稳健刻画时序结构
+- 因果 mask 与 padding 感知归一化贴合真实日志
+- 绑权输出头降低参数量并提升训练稳定性
+- 直接作为生成式推荐的骨干网络
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from nextrec.basic.model import BaseModel
+from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
+
+
+def _relative_position_bucket(
+    relative_position: torch.Tensor,
+    num_buckets: int = 32,
+    max_distance: int = 128,
+) -> torch.Tensor:
+    """
+    map the relative position (i-j) to a bucket in [0, num_buckets).
+    """
+    # only need the negative part for causal attention
+    n = -relative_position
+    n = torch.clamp(n, min=0)
+
+    # when the distance is small, keep it exact
+    max_exact = num_buckets // 2
+    is_small = n < max_exact
+
+    # when the distance is too far, do log scaling
+    large_val = max_exact + ((torch.log(n.float() / max_exact + 1e-6) / math.log(max_distance / max_exact)) * (num_buckets - max_exact)).long()
+    large_val = torch.clamp(large_val, max=num_buckets - 1)
+
+    buckets = torch.where(is_small, n.long(), large_val)
+    return buckets
+
+
+class RelativePositionBias(nn.Module):
+    """
+    Compute relative position bias (RAB) for HSTU attention.
+    The input is the sequence length T, output is [1, num_heads, seq_len, seq_len].
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        num_buckets: int = 32,
+        max_distance: int = 128,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.embedding = nn.Embedding(num_buckets, num_heads)
+
+    def forward(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        # positions: [T]
+        ctx = torch.arange(seq_len, device=device)[:, None]
+        mem = torch.arange(seq_len, device=device)[None, :]
+        rel_pos = mem - ctx  # a matrix to describe all relative positions for each [i,j] pair, shape = [seq_len, seq_len]
+        buckets = _relative_position_bucket(rel_pos, num_buckets=self.num_buckets, max_distance=self.max_distance,)  # map to buckets
+        values = self.embedding(buckets) # embedding vector for each [i,j] pair, shape = [seq_len, seq_len, embedding_dim=num_heads]
+        return values.permute(2, 0, 1).unsqueeze(0) # [1, num_heads, seq_len, seq_len]
+
+class HSTUPointwiseAttention(nn.Module):
+    """
+    Pointwise aggregation attention that implements HSTU without softmax:
+        1) [U, V, Q, K] = split( φ1(f1(X)) ), U: gate, V: value, Q: query, K: key
+        2) AV = φ2(QK^T + rab) V / N, av is attention-weighted value
+        3) Y  = f2( Norm(AV) ⊙ U ), y is output
+    φ1, φ2 use SiLU; Norm uses LayerNorm.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dropout: float = 0.1,
+        alpha: float | None = None
+    ):
+        super().__init__()
+        if d_model % num_heads != 0:
+            raise ValueError(f"[HSTUPointwiseAttention Error] d_model({d_model}) % num_heads({num_heads}) != 0")
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.d_head = d_model // num_heads
+        self.alpha = alpha if alpha is not None else (self.d_head ** -0.5)
+        # project input to 4 * d_model for U, V, Q, K
+        self.in_proj = nn.Linear(d_model, 4 * d_model, bias=True)
+        # project output back to d_model
+        self.out_proj = nn.Linear(d_model, d_model, bias=True)
+        self.dropout = nn.Dropout(dropout)
+        self.norm = nn.LayerNorm(d_model)
+
+    def _reshape_heads(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        [B, T, D] -> [B, H, T, d_head]
+        """
+        B, T, D = x.shape
+        return x.view(B, T, self.num_heads, self.d_head).transpose(1, 2)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,      # [T, T] with 0 or -inf
+        key_padding_mask: Optional[torch.Tensor] = None,  # [B, T], True = pad
+        rab: Optional[torch.Tensor] = None,            # [1, H, T, T] or None
+    ) -> torch.Tensor:
+        B, T, D = x.shape
+
+        # Eq.(1): [U, V, Q, K] = split( φ1(f1(X)) )
+        h = F.silu(self.in_proj(x))  # [B, T, 4D]
+        U, V, Q, K = h.chunk(4, dim=-1)  # each [B, T, D]
+
+        Qh = self._reshape_heads(Q)  # [B, H, T, d_head]
+        Kh = self._reshape_heads(K)  # [B, H, T, d_head]
+        Vh = self._reshape_heads(V)  # [B, H, T, d_head]
+        Uh = self._reshape_heads(U)  # [B, H, T, d_head]
+
+        # attention logits: QK^T (without 1/sqrt(d) and softmax)
+        logits = torch.matmul(Qh, Kh.transpose(-2, -1)) * self.alpha  # [B, H, T, T]
+
+        # add relative position bias (rab^p), and future extensible rab^t
+        if rab is not None:
+            # rab: [1, H, T, T] or [B, H, T, T]
+            logits = logits + rab
+
+        # construct an "allowed" mask to calculate N
+        # 1 indicates that the (query i, key j) pair is a valid attention pair; 0 indicates it is masked out
+        allowed = torch.ones_like(logits, dtype=torch.float)  # [B, H, T, T]
+
+        # causal mask: attn_mask is usually an upper triangular matrix of -inf with shape [T, T]
+        if attn_mask is not None:
+            allowed = allowed * (attn_mask.view(1, 1, T, T) == 0).float()
+            logits = logits + attn_mask.view(1, 1, T, T)
+
+        # padding mask: key_padding_mask is usually [B, T], True = pad
+        if key_padding_mask is not None:
+            # valid: 1 for non-pad, 0 for pad
+            valid = (~key_padding_mask).float()          # [B, T]
+            valid = valid.view(B, 1, 1, T)               # [B, 1, 1, T]
+            allowed = allowed * valid
+            logits = logits.masked_fill(valid == 0, float("-inf"))
+
+        # Eq.(2): A(X)V(X) = φ2(QK^T + rab) V(X) / N
+        attn = F.silu(logits)  # [B, H, T, T]
+        denom = allowed.sum(dim=-1, keepdim=True)  # [B, H, T, 1]
+        denom = denom.clamp(min=1.0)
+
+        attn = attn / denom  # [B, H, T, T]
+        AV = torch.matmul(attn, Vh)  # [B, H, T, d_head]
+        AV = AV.transpose(1, 2).contiguous().view(B, T, D) # reshape back to [B, T, D]
+        U_flat = Uh.transpose(1, 2).contiguous().view(B, T, D)
+        y = self.out_proj(self.dropout(self.norm(AV) * U_flat))  # [B, T, D]
+        return y
+
+
+class HSTULayer(nn.Module):
+    """
+    HSTUPointwiseAttention + Residual Connection
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dropout: float = 0.1,
+        use_rab_pos: bool = True,
+        rab_num_buckets: int = 32,
+        rab_max_distance: int = 128,
+    ):
+        super().__init__()
+        self.attn = HSTUPointwiseAttention(d_model=d_model, num_heads=num_heads, dropout=dropout)
+        self.dropout = nn.Dropout(dropout)
+        self.use_rab_pos = use_rab_pos
+        self.rel_pos_bias = (RelativePositionBias(num_heads=num_heads, num_buckets=rab_num_buckets, max_distance=rab_max_distance) if use_rab_pos else None)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        key_padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        x: [B, T, D]
+        """
+        B, T, D = x.shape
+        device = x.device
+        rab = None
+        if self.use_rab_pos:
+            rab = self.rel_pos_bias(seq_len=T, device=device) # [1, H, T, T]
+        out = self.attn(x=x, attn_mask=attn_mask, key_padding_mask=key_padding_mask, rab=rab)
+        return x + self.dropout(out)
+
+
+class HSTU(BaseModel):
+    """
+    HSTU encoder for next-item prediction in a causal, generative setup.
+    Pipeline:
+      1) Embed tokens + positions from the behavior history
+      2) Apply stacked HSTU layers with causal mask and optional RAB
+      3) Use the last valid position to produce next-item logits via tied LM head
+    """
+
+    @property
+    def model_name(self) -> str:
+        return "HSTU"
+
+    @property
+    def task_type(self) -> str:
+        return "multiclass"
+
+    def __init__(
+        self,
+        sequence_features: list[SequenceFeature],
+        dense_features: Optional[list[DenseFeature]] = None,
+        sparse_features: Optional[list[SparseFeature]] = None,
+        d_model: Optional[int] = None,
+        num_heads: int = 8,
+        num_layers: int = 4,
+        max_seq_len: int = 200,
+        dropout: float = 0.1,
+        # RAB settings
+        use_rab_pos: bool = True,
+        rab_num_buckets: int = 32,
+        rab_max_distance: int = 128,
+
+        tie_embeddings: bool = True,
+        target: Optional[list[str] | str] = None,
+        optimizer: str = "adam",
+        optimizer_params: Optional[dict] = None,
+        scheduler: Optional[str] = None,
+        scheduler_params: Optional[dict] = None,
+        loss_params: Optional[dict] = None,
+        embedding_l1_reg: float = 0.0,
+        dense_l1_reg: float = 0.0,
+        embedding_l2_reg: float = 0.0,
+        dense_l2_reg: float = 0.0,
+        device: str = "cpu",
+        **kwargs,
+    ):
+        if not sequence_features:
+            raise ValueError("[HSTU Error] HSTU requires at least one SequenceFeature (user behavior history).")
+
+        # demo version: use the first SequenceFeature as the main sequence
+        self.history_feature = sequence_features[0]
+
+        hidden_dim = d_model or max(int(getattr(self.history_feature, "embedding_dim", 0) or 0), 32)
+        # Make hidden_dim divisible by num_heads
+        if hidden_dim % num_heads != 0:
+            hidden_dim = num_heads * math.ceil(hidden_dim / num_heads)
+
+        self.padding_idx = self.history_feature.padding_idx if self.history_feature.padding_idx is not None else 0
+        self.vocab_size = self.history_feature.vocab_size
+        self.max_seq_len = max_seq_len
+
+        super().__init__(
+            dense_features=dense_features,
+            sparse_features=sparse_features,
+            sequence_features=sequence_features,
+            target=target,
+            task=self.task_type,
+            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,
+            **kwargs,
+        )
+
+        # token & position embedding (paper usually includes pos embedding / RAB in encoder)
+        self.token_embedding = nn.Embedding(
+            num_embeddings=self.vocab_size,
+            embedding_dim=hidden_dim,
+            padding_idx=self.padding_idx,
+        )
+        self.position_embedding = nn.Embedding(max_seq_len, hidden_dim)
+        self.input_dropout = nn.Dropout(dropout)
+
+        # HSTU layers
+        self.layers = nn.ModuleList([HSTULayer(d_model=hidden_dim, num_heads=num_heads, dropout=dropout, use_rab_pos=use_rab_pos, 
+                                               rab_num_buckets=rab_num_buckets, rab_max_distance=rab_max_distance) for _ in range(num_layers)])
+
+        self.final_norm = nn.LayerNorm(hidden_dim)
+        self.lm_head = nn.Linear(hidden_dim, self.vocab_size, bias=False)
+        if tie_embeddings:
+            self.lm_head.weight = self.token_embedding.weight
+
+        # causal mask buffer
+        self.register_buffer("causal_mask", torch.empty(0), persistent=False)
+        self.ignore_index = self.padding_idx if self.padding_idx is not None else -100
+
+        optimizer_params = optimizer_params or {}
+        scheduler_params = scheduler_params or {}
+        loss_params = loss_params or {}
+        loss_params.setdefault("ignore_index", self.ignore_index)
+
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, scheduler=scheduler, scheduler_params=scheduler_params, loss="crossentropy", loss_params=loss_params)
+        self._register_regularization_weights(embedding_attr="token_embedding", include_modules=["layers", "lm_head"])
+
+    def _build_causal_mask(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        """
+        build causal mask of shape [T, T]: upper triangle is -inf, others are 0.
+        This will be added to the logits to simulate causal structure.
+        """
+        if self.causal_mask.numel() == 0 or self.causal_mask.size(0) < seq_len:
+            mask = torch.full((seq_len, seq_len), float("-inf"), device=device)
+            mask = torch.triu(mask, diagonal=1) 
+            self.causal_mask = mask
+        return self.causal_mask[:seq_len, :seq_len]
+
+    def _trim_sequence(self, seq: torch.Tensor) -> torch.Tensor:
+        if seq.size(1) <= self.max_seq_len:
+            return seq
+        return seq[:, -self.max_seq_len :]
+
+    def forward(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
+        seq = x[self.history_feature.name].long()  # [B, T_raw]
+        seq = self._trim_sequence(seq)            # [B, T]
+
+        B, T = seq.shape
+        device = seq.device
+        # position ids: [B, T]
+        pos_ids = torch.arange(T, device=device).unsqueeze(0).expand(B, -1)
+        token_emb = self.token_embedding(seq)         # [B, T, D]
+        pos_emb = self.position_embedding(pos_ids)    # [B, T, D]
+        hidden_states = self.input_dropout(token_emb + pos_emb)
+
+        # padding mask：True = pad
+        padding_mask = seq.eq(self.padding_idx)       # [B, T]
+        attn_mask = self._build_causal_mask(seq_len=T, device=device)  # [T, T]
+
+        for layer in self.layers:
+            hidden_states = layer(x=hidden_states, attn_mask=attn_mask, key_padding_mask=padding_mask)
+        hidden_states = self.final_norm(hidden_states)  # [B, T, D]
+
+        valid_lengths = (~padding_mask).sum(dim=1)  # [B]
+        last_index = (valid_lengths - 1).clamp(min=0)
+        last_hidden = hidden_states[torch.arange(B, device=device), last_index]  # [B, D]
+
+        logits = self.lm_head(last_hidden)  # [B, vocab_size]
+        return logits
+
+    def compute_loss(self, y_pred, y_true):
+        """
+        y_true: [B] or [B, 1], the id of the next item.
+        """
+        if y_true is None:
+            raise ValueError("[HSTU-compute_loss] Training requires y_true (next item id).")
+        labels = y_true.view(-1).long()
+        return self.loss_fn[0](y_pred, labels)

nextrec/models/match/mind.py

@@ -13,7 +13,116 @@ from typing import Literal
 
 from nextrec.basic.model import BaseMatchModel
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, CapsuleNetwork
+from nextrec.basic.layers import MLP, EmbeddingLayer
+
+class MultiInterestSA(nn.Module):
+    """Multi-interest self-attention extractor from MIND (Li et al., 2019)."""
+
+    def __init__(self, embedding_dim, interest_num, hidden_dim=None):
+        super(MultiInterestSA, self).__init__()
+        self.embedding_dim = embedding_dim
+        self.interest_num = interest_num
+        if hidden_dim == None:
+            self.hidden_dim = self.embedding_dim * 4
+        self.W1 = torch.nn.Parameter(torch.rand(self.embedding_dim, self.hidden_dim), requires_grad=True)
+        self.W2 = torch.nn.Parameter(torch.rand(self.hidden_dim, self.interest_num), requires_grad=True)
+        self.W3 = torch.nn.Parameter(torch.rand(self.embedding_dim, self.embedding_dim), requires_grad=True)
+
+    def forward(self, seq_emb, mask=None):
+        H = torch.einsum('bse, ed -> bsd', seq_emb, self.W1).tanh()
+        if mask != None:
+            A = torch.einsum('bsd, dk -> bsk', H, self.W2) + -1.e9 * (1 - mask.float())
+            A = F.softmax(A, dim=1)
+        else:
+            A = F.softmax(torch.einsum('bsd, dk -> bsk', H, self.W2), dim=1)
+        A = A.permute(0, 2, 1)
+        multi_interest_emb = torch.matmul(A, seq_emb)
+        return multi_interest_emb
+
+
+class CapsuleNetwork(nn.Module):
+    """Dynamic routing capsule network used in MIND (Li et al., 2019)."""
+
+    def __init__(self, embedding_dim, seq_len, bilinear_type=2, interest_num=4, routing_times=3, relu_layer=False):
+        super(CapsuleNetwork, self).__init__()
+        self.embedding_dim = embedding_dim  # h
+        self.seq_len = seq_len  # s
+        self.bilinear_type = bilinear_type
+        self.interest_num = interest_num
+        self.routing_times = routing_times
+
+        self.relu_layer = relu_layer
+        self.stop_grad = True
+        self.relu = nn.Sequential(nn.Linear(self.embedding_dim, self.embedding_dim, bias=False), nn.ReLU())
+        if self.bilinear_type == 0:  # MIND
+            self.linear = nn.Linear(self.embedding_dim, self.embedding_dim, bias=False)
+        elif self.bilinear_type == 1:
+            self.linear = nn.Linear(self.embedding_dim, self.embedding_dim * self.interest_num, bias=False)
+        else:
+            self.w = nn.Parameter(torch.Tensor(1, self.seq_len, self.interest_num * self.embedding_dim, self.embedding_dim))
+            nn.init.xavier_uniform_(self.w)
+
+    def forward(self, item_eb, mask):
+        if self.bilinear_type == 0:
+            item_eb_hat = self.linear(item_eb)
+            item_eb_hat = item_eb_hat.repeat(1, 1, self.interest_num)
+        elif self.bilinear_type == 1:
+            item_eb_hat = self.linear(item_eb)
+        else:
+            u = torch.unsqueeze(item_eb, dim=2)
+            item_eb_hat = torch.sum(self.w[:, :self.seq_len, :, :] * u, dim=3)
+
+        item_eb_hat = torch.reshape(item_eb_hat, (-1, self.seq_len, self.interest_num, self.embedding_dim))
+        item_eb_hat = torch.transpose(item_eb_hat, 1, 2).contiguous()
+        item_eb_hat = torch.reshape(item_eb_hat, (-1, self.interest_num, self.seq_len, self.embedding_dim))
+
+        if self.stop_grad:
+            item_eb_hat_iter = item_eb_hat.detach()
+        else:
+            item_eb_hat_iter = item_eb_hat
+
+        if self.bilinear_type > 0:
+            capsule_weight = torch.zeros(item_eb_hat.shape[0],
+                                         self.interest_num,
+                                         self.seq_len,
+                                         device=item_eb.device,
+                                         requires_grad=False)
+        else:
+            capsule_weight = torch.randn(item_eb_hat.shape[0],
+                                         self.interest_num,
+                                         self.seq_len,
+                                         device=item_eb.device,
+                                         requires_grad=False)
+
+        for i in range(self.routing_times):  # 动态路由传播3次
+            atten_mask = torch.unsqueeze(mask, 1).repeat(1, self.interest_num, 1)
+            paddings = torch.zeros_like(atten_mask, dtype=torch.float)
+
+            capsule_softmax_weight = F.softmax(capsule_weight, dim=-1)
+            capsule_softmax_weight = torch.where(torch.eq(atten_mask, 0), paddings, capsule_softmax_weight)
+            capsule_softmax_weight = torch.unsqueeze(capsule_softmax_weight, 2)
+
+            if i < 2:
+                interest_capsule = torch.matmul(capsule_softmax_weight, item_eb_hat_iter)
+                cap_norm = torch.sum(torch.square(interest_capsule), -1, True)
+                scalar_factor = cap_norm / (1 + cap_norm) / torch.sqrt(cap_norm + 1e-9)
+                interest_capsule = scalar_factor * interest_capsule
+
+                delta_weight = torch.matmul(item_eb_hat_iter, torch.transpose(interest_capsule, 2, 3).contiguous())
+                delta_weight = torch.reshape(delta_weight, (-1, self.interest_num, self.seq_len))
+                capsule_weight = capsule_weight + delta_weight
+            else:
+                interest_capsule = torch.matmul(capsule_softmax_weight, item_eb_hat)
+                cap_norm = torch.sum(torch.square(interest_capsule), -1, True)
+                scalar_factor = cap_norm / (1 + cap_norm) / torch.sqrt(cap_norm + 1e-9)
+                interest_capsule = scalar_factor * interest_capsule
+
+        interest_capsule = torch.reshape(interest_capsule, (-1, self.interest_num, self.embedding_dim))
+
+        if self.relu_layer:
+            interest_capsule = self.relu(interest_capsule)
+
+        return interest_capsule
 
 
 class MIND(BaseMatchModel):

nextrec/models/multi_task/esmm.py

@@ -1,7 +1,44 @@
 """
 Date: create on 09/11/2025
+Checkpoint: edit on 29/11/2025
 Author: Yang Zhou,zyaztec@gmail.com
-Reference: [1] Ma X, Zhao L, Huang G, et al. Entire space multi-task model: An effective approach for estimating post-click conversion rate[C]//SIGIR. 2018: 1137-1140.
+Reference:
+[1] Ma X, Zhao L, Huang G, et al. Entire space multi-task model: An effective approach
+for estimating post-click conversion rate[C]//SIGIR. 2018: 1137-1140.
+(https://dl.acm.org/doi/10.1145/3209978.3210007)
+
+Entire Space Multi-task Model (ESMM) targets CVR estimation by jointly optimizing
+CTR and CTCVR on the full impression space, mitigating sample selection bias and
+conversion sparsity. CTR predicts P(click | impression), CVR predicts P(conversion |
+click), and their product forms CTCVR supervised on impression labels.
+
+Workflow:
+  (1) Shared embeddings encode all features from impressions
+  (2) CTR tower outputs click probability conditioned on impression
+  (3) CVR tower outputs conversion probability conditioned on click
+  (4) CTCVR = CTR * CVR enables end-to-end training without filtering clicked data
+
+Key Advantages:
+- Trains on the entire impression space to remove selection bias
+- Transfers rich click signals to sparse conversion prediction via shared embeddings
+- Stable optimization by decomposing CTCVR into well-defined sub-tasks
+- Simple architecture that can pair with other multi-task variants
+
+ESMM（Entire Space Multi-task Model）用于 CVR 预估，通过在曝光全空间联合训练
+CTR 与 CTCVR，缓解样本选择偏差和转化数据稀疏问题。CTR 预测 P(click|impression)，
+CVR 预测 P(conversion|click)，二者相乘得到 CTCVR 并在曝光标签上直接监督。
+
+流程：
+  (1) 共享 embedding 统一处理曝光特征
+  (2) CTR 塔输出曝光下的点击概率
+  (3) CVR 塔输出点击后的转化概率
+  (4) CTR 与 CVR 相乘得到 CTCVR，无需只在点击子集上训练
+
+主要优点：
+- 在曝光空间训练，避免样本选择偏差
+- 通过共享表示将点击信号迁移到稀疏的转化任务
+- 将 CTCVR 分解为子任务，优化稳定
+- 结构简单，可与其它多任务方法组合使用
 """
 
 import torch
@@ -77,37 +114,22 @@ class ESMM(BaseModel):
             
         # All features
         self.all_features = dense_features + sparse_features + sequence_features
-
         # Shared embedding layer
         self.embedding = EmbeddingLayer(features=self.all_features)
+        input_dim = self.embedding.input_dim # Calculate input dimension, better way than below
+        # 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
 
-        # 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
-        
         # CTR tower
         self.ctr_tower = MLP(input_dim=input_dim, output_layer=True, **ctr_params)
         
         # CVR tower
         self.cvr_tower = MLP(input_dim=input_dim, output_layer=True, **cvr_params)
-        self.prediction_layer = PredictionLayer(
-            task_type=self.task_type,
-            task_dims=[1, 1]
-        )
-
+        self.prediction_layer = PredictionLayer(task_type=self.task_type, task_dims=[1, 1])
         # Register regularization weights
-        self._register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=['ctr_tower', 'cvr_tower']
-        )
-
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss,
-            loss_params=loss_params,
-        )
+        self._register_regularization_weights(embedding_attr='embedding', include_modules=['ctr_tower', 'cvr_tower'])
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)
 
     def forward(self, x):
         # Get all embeddings and flatten
@@ -119,11 +141,8 @@ class ESMM(BaseModel):
         logits = torch.cat([ctr_logit, cvr_logit], dim=1)
         preds = self.prediction_layer(logits)
         ctr, cvr = preds.chunk(2, dim=1)
-        
-        # CTCVR prediction: P(click & conversion | impression) = P(click) * P(conversion | click)
         ctcvr = ctr * cvr  # [B, 1]
         
-        # Output: [CTR, CTCVR]
-        # Note: We supervise CTR with click labels and CTCVR with conversion labels
+        # Output: [CTR, CTCVR], We supervise CTR with click labels and CTCVR with conversion labels
         y = torch.cat([ctr, ctcvr], dim=1)  # [B, 2]
         return y  # [B, 2], where y[:, 0] is CTR and y[:, 1] is CTCVR