nextrec 0.4.11__py3-none-any.whl → 0.4.13__py3-none-any.whl

nextrec/loss/grad_norm.py

@@ -0,0 +1,232 @@
+"""
+GradNorm loss weighting for multi-task learning.
+
+Date: create on 27/10/2025
+Checkpoint: edit on 20/12/2025
+Author: Yang Zhou,zyaztec@gmail.com
+
+Reference:
+Chen, Zhao, et al. "GradNorm: Gradient Normalization for Adaptive Loss Balancing
+in Deep Multitask Networks." ICML 2018.
+"""
+
+from __future__ import annotations
+
+from typing import Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def get_grad_norm_shared_params(
+    model: torch.nn.Module,
+    shared_modules: Iterable[str] | None = None,
+) -> list[torch.nn.Parameter]:
+    if not shared_modules:
+        return [p for p in model.parameters() if p.requires_grad]
+    shared_params = []
+    seen = set()
+    for name in shared_modules:
+        module = getattr(model, name, None)
+        if module is None:
+            continue
+        for param in module.parameters():
+            if param.requires_grad and id(param) not in seen:
+                shared_params.append(param)
+                seen.add(id(param))
+    if not shared_params:
+        return [p for p in model.parameters() if p.requires_grad]
+    return shared_params
+
+
+class GradNormLossWeighting:
+    """
+    Adaptive multi-task loss weighting with GradNorm.
+
+    Args:
+        num_tasks: Number of tasks.
+        alpha: GradNorm balancing strength.
+        lr: Learning rate for the weight optimizer.
+        init_weights: Optional initial weights per task.
+        device: Torch device for weights.
+        ema_decay: Optional EMA decay for smoothing loss ratios.
+        init_ema_steps: Number of steps to build EMA for initial losses.
+        init_ema_decay: EMA decay for initial losses when init_ema_steps > 0.
+        eps: Small value for numerical stability.
+    """
+
+    def __init__(
+        self,
+        num_tasks: int,
+        alpha: float = 1.5,
+        lr: float = 0.025,
+        init_weights: Iterable[float] | None = None,
+        device: torch.device | str | None = None,
+        ema_decay: float | None = None,
+        init_ema_steps: int = 0,
+        init_ema_decay: float = 0.9,
+        eps: float = 1e-8,
+    ) -> None:
+        if num_tasks <= 1:
+            raise ValueError("GradNorm requires num_tasks > 1.")
+        self.num_tasks = num_tasks
+        self.alpha = alpha
+        self.eps = eps
+        if ema_decay is not None:
+            ema_decay = ema_decay
+            if ema_decay < 0.0 or ema_decay >= 1.0:
+                raise ValueError("ema_decay must be in [0.0, 1.0).")
+        self.ema_decay = ema_decay
+        self.init_ema_steps = init_ema_steps
+        if self.init_ema_steps < 0:
+            raise ValueError("init_ema_steps must be >= 0.")
+        self.init_ema_decay = init_ema_decay
+        if self.init_ema_decay < 0.0 or self.init_ema_decay >= 1.0:
+            raise ValueError("init_ema_decay must be in [0.0, 1.0).")
+        self.init_ema_count = 0
+
+        if init_weights is None:
+            weights = torch.ones(self.num_tasks, dtype=torch.float32)
+        else:
+            weights = torch.tensor(list(init_weights), dtype=torch.float32)
+            if weights.numel() != self.num_tasks:
+                raise ValueError(
+                    "init_weights length must match num_tasks for GradNorm."
+                )
+        if device is not None:
+            weights = weights.to(device)
+        self.weights = nn.Parameter(weights)
+        self.optimizer = torch.optim.Adam([self.weights], lr=float(lr))
+
+        self.initial_losses = None
+        self.initial_losses_ema = None
+        self.loss_ema = None
+        self.pending_grad = None
+
+    def to(self, device):
+        device = torch.device(device)
+        self.weights.data = self.weights.data.to(device)
+        if self.initial_losses is not None:
+            self.initial_losses = self.initial_losses.to(device)
+        if self.initial_losses_ema is not None:
+            self.initial_losses_ema = self.initial_losses_ema.to(device)
+        if self.loss_ema is not None:
+            self.loss_ema = self.loss_ema.to(device)
+        return self
+
+    def compute_weighted_loss(
+        self,
+        task_losses: list[torch.Tensor],
+        shared_params: Iterable[torch.nn.Parameter],
+    ) -> torch.Tensor:
+        """
+        Return weighted total loss and update task weights with GradNorm.
+        """
+        if len(task_losses) != self.num_tasks:
+            raise ValueError(
+                f"Expected {self.num_tasks} task losses, got {len(task_losses)}."
+            )
+        shared_params = [p for p in shared_params if p.requires_grad]
+        if not shared_params:
+            return torch.stack(task_losses).sum()
+
+        with torch.no_grad():
+            loss_values = torch.tensor(
+                [loss.item() for loss in task_losses], device=self.weights.device
+            )
+            if self.initial_losses is None:
+                if self.init_ema_steps > 0:
+                    if self.initial_losses_ema is None:
+                        self.initial_losses_ema = loss_values
+                    else:
+                        self.initial_losses_ema = (
+                            self.init_ema_decay * self.initial_losses_ema
+                            + (1.0 - self.init_ema_decay) * loss_values
+                        )
+                    self.init_ema_count += 1
+                    if self.init_ema_count >= self.init_ema_steps:
+                        self.initial_losses = self.initial_losses_ema.clone()
+                else:
+                    self.initial_losses = loss_values
+
+        weights_detached = self.weights.detach()
+        weighted_losses = [
+            weights_detached[i] * task_losses[i] for i in range(self.num_tasks)
+        ]
+        total_loss = torch.stack(weighted_losses).sum()
+
+        grad_norms = self.compute_grad_norms(task_losses, shared_params)
+        with torch.no_grad():
+            if self.ema_decay is not None:
+                if self.loss_ema is None:
+                    self.loss_ema = loss_values
+                else:
+                    self.loss_ema = (
+                        self.ema_decay * self.loss_ema
+                        + (1.0 - self.ema_decay) * loss_values
+                    )
+                ratio_source = self.loss_ema
+            else:
+                ratio_source = loss_values
+            if self.initial_losses is not None:
+                base_initial = self.initial_losses
+            elif self.initial_losses_ema is not None:
+                base_initial = self.initial_losses_ema
+            else:
+                base_initial = loss_values
+            loss_ratios = ratio_source / (base_initial + self.eps)
+            inv_rate = loss_ratios / (loss_ratios.mean() + self.eps)
+            target = grad_norms.mean() * (inv_rate**self.alpha)
+
+        grad_norm_loss = F.l1_loss(grad_norms, target.detach(), reduction="sum")
+        grad_w = torch.autograd.grad(grad_norm_loss, self.weights, retain_graph=True)[0]
+        self.pending_grad = grad_w.detach()
+
+        return total_loss
+
+    def compute_grad_norms(self, task_losses, shared_params):
+        grad_norms = []
+        for i, task_loss in enumerate(task_losses):
+            grads = torch.autograd.grad(
+                self.weights[i] * task_loss,
+                shared_params,
+                retain_graph=True,
+                create_graph=True,
+                allow_unused=True,
+            )
+
+            sq_sum = torch.zeros((), device=self.weights.device)
+            any_used = False
+            for g in grads:
+                if g is not None:
+                    any_used = True
+                    sq_sum = sq_sum + g.pow(2).sum()
+
+            if not any_used:
+                total_norm = torch.tensor(self.eps, device=self.weights.device)
+            else:
+                total_norm = torch.sqrt(sq_sum + self.eps)
+
+            grad_norms.append(total_norm)
+
+        return torch.stack(grad_norms)
+
+    def step(self) -> None:
+        if self.pending_grad is None:
+            return
+
+        self.optimizer.zero_grad(set_to_none=True)
+
+        if self.weights.grad is None:
+            self.weights.grad = torch.zeros_like(self.weights)
+        self.weights.grad.copy_(self.pending_grad)
+
+        self.optimizer.step()
+
+        with torch.no_grad():
+            w = self.weights.clamp(min=self.eps)
+            w = w * self.num_tasks / (w.sum() + self.eps)
+            self.weights.copy_(w)
+
+        self.pending_grad = None

nextrec/loss/loss_utils.py

@@ -60,7 +60,7 @@ def build_cb_focal(kw):
     return ClassBalancedFocalLoss(**kw)
 
 
-def get_loss_fn(loss: LossType | nn.Module | None = None, **kw) -> nn.Module:
+def get_loss_fn(loss=None, **kw) -> nn.Module:
     """
     Get loss function by name or return the provided loss module.
 

nextrec/models/multi_task/esmm.py

@@ -138,6 +138,7 @@ class ESMM(BaseModel):
 
         # CVR tower
         self.cvr_tower = MLP(input_dim=input_dim, output_layer=True, **cvr_params)
+        self.grad_norm_shared_modules = ["embedding"]
         self.prediction_layer = PredictionLayer(
             task_type=self.default_task, task_dims=[1, 1]
         )

nextrec/models/multi_task/mmoe.py

@@ -165,6 +165,7 @@ class MMOE(BaseModel):
         for _ in range(self.num_tasks):
             gate = nn.Sequential(nn.Linear(input_dim, num_experts), nn.Softmax(dim=1))
             self.gates.append(gate)
+        self.grad_norm_shared_modules = ["embedding", "experts", "gates"]
 
         # Task-specific towers
         self.towers = nn.ModuleList()

nextrec/models/multi_task/ple.py

@@ -295,6 +295,7 @@ class PLE(BaseModel):
             )
             self.cgc_layers.append(cgc_layer)
             expert_output_dim = cgc_layer.output_dim
+        self.grad_norm_shared_modules = ["embedding", "cgc_layers"]
 
         # Task-specific towers
         self.towers = nn.ModuleList()

nextrec/models/multi_task/poso.py

@@ -483,6 +483,10 @@ class POSO(BaseModel):
                 ]
             )
             self.tower_heads = None
+        if self.architecture == "mlp":
+            self.grad_norm_shared_modules = ["embedding"]
+        else:
+            self.grad_norm_shared_modules = ["embedding", "mmoe"]
         self.prediction_layer = PredictionLayer(
             task_type=self.default_task,
             task_dims=[1] * self.num_tasks,

nextrec/models/multi_task/share_bottom.py

@@ -129,6 +129,7 @@ class ShareBottom(BaseModel):
 
         # Shared bottom network
         self.bottom = MLP(input_dim=input_dim, output_layer=False, **bottom_params)
+        self.grad_norm_shared_modules = ["embedding", "bottom"]
 
         # Get bottom output dimension
         if "dims" in bottom_params and len(bottom_params["dims"]) > 0:

nextrec/models/ranking/eulernet.py

@@ -41,7 +41,8 @@ from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import LR, EmbeddingLayer, PredictionLayer
 from nextrec.basic.model import BaseModel
 
-class EulerInteractionLayerPaper(nn.Module):
+
+class EulerInteractionLayer(nn.Module):
     """
     Paper-aligned Euler Interaction Layer.
 
@@ -102,24 +103,32 @@ class EulerInteractionLayerPaper(nn.Module):
             self.bn = None
             self.ln = None
 
-    def forward(self, r: torch.Tensor, p: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def forward(
+        self, r: torch.Tensor, p: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         r, p: [B, m, d]
         return r_out, p_out: [B, n, d]
         """
         B, m, d = r.shape
-        assert m == self.m and d == self.d, f"Expected [B,{self.m},{self.d}] got {r.shape}"
+        assert (
+            m == self.m and d == self.d
+        ), f"Expected [B,{self.m},{self.d}] got {r.shape}"
 
         # Euler Transformation: rectangular -> polar
-        lam = torch.sqrt(r * r + p * p + self.eps)         # [B,m,d]
-        theta = torch.atan2(p, r)              # [B,m,d]
-        log_lam = torch.log(lam + self.eps)               # [B,m,d]
+        lam = torch.sqrt(r * r + p * p + self.eps)  # [B,m,d]
+        theta = torch.atan2(p, r)  # [B,m,d]
+        log_lam = torch.log(lam + self.eps)  # [B,m,d]
 
         # Generalized Multi-order Transformation
         # psi_k = sum_j alpha_{k,j} * theta_j + delta_k
         # l_k   = exp(sum_j alpha_{k,j} * log(lam_j) + delta'_k)
-        psi = torch.einsum("bmd,nmd->bnd", theta, self.alpha) + self.delta_phase  # [B,n,d]
-        log_l = torch.einsum("bmd,nmd->bnd", log_lam, self.alpha) + self.delta_logmod  # [B,n,d]
+        psi = (
+            torch.einsum("bmd,nmd->bnd", theta, self.alpha) + self.delta_phase
+        )  # [B,n,d]
+        log_l = (
+            torch.einsum("bmd,nmd->bnd", log_lam, self.alpha) + self.delta_logmod
+        )  # [B,n,d]
         l = torch.exp(log_l)  # [B,n,d]
 
         # Inverse Euler Transformation
@@ -153,7 +162,7 @@ class EulerInteractionLayerPaper(nn.Module):
         return r_out, p_out
 
 
-class ComplexSpaceMappingPaper(nn.Module):
+class ComplexSpaceMapping(nn.Module):
     """
     Map real embeddings e_j to complex features via Euler's formula (Eq.6-7).
     For each field j:
@@ -174,63 +183,6 @@ class ComplexSpaceMappingPaper(nn.Module):
         r = mu * torch.cos(e)
         p = mu * torch.sin(e)
         return r, p
-    
-class EulerNetPaper(nn.Module):
-    """
-    Paper-aligned EulerNet core (embedding -> mapping -> L Euler layers -> linear regression).
-    """
-
-    def __init__(
-        self,
-        *,
-        embedding_dim: int,
-        num_fields: int,
-        num_layers: int = 2,
-        num_orders: int = 8,      # n in paper
-        use_implicit: bool = True,
-        norm: str | None = "ln",  # None | "bn" | "ln"
-    ):
-        super().__init__()
-        self.d = embedding_dim
-        self.m = num_fields
-        self.L = num_layers
-        self.n = num_orders
-
-        self.mapping = ComplexSpaceMappingPaper(embedding_dim, num_fields)
-
-        self.layers = nn.ModuleList([
-            EulerInteractionLayerPaper(
-                embedding_dim=embedding_dim,
-                num_fields=(num_fields if i == 0 else num_orders),  # stack: m -> n -> n ...
-                num_orders=num_orders,
-                use_implicit=use_implicit,
-                norm=norm,
-            )
-            for i in range(num_layers)
-        ])
-
-        # Output regression (Eq.16-17)
-        # After last layer: r,p are [B,n,d]. Concatenate to [B, n*d] each, then regress.
-        self.w = nn.Linear(self.n * self.d, 1, bias=False)  # for real
-        self.w_im = nn.Linear(self.n * self.d, 1, bias=False)  # for imag
-
-    def forward(self, field_emb: torch.Tensor) -> torch.Tensor:
-        """
-        field_emb: [B, m, d] real embeddings e_j
-        return: logits, shape [B,1]
-        """
-        r, p = self.mapping(field_emb)  # [B,m,d]
-
-        # stack Euler interaction layers
-        for layer in self.layers:
-            r, p = layer(r, p)  # -> [B,n,d]
-
-        r_flat = r.reshape(r.size(0), self.n * self.d)
-        p_flat = p.reshape(p.size(0), self.n * self.d)
-
-        z_re = self.w(r_flat)
-        z_im = self.w_im(p_flat)
-        return z_re + z_im  # Eq.17 logits
 
 
 class EulerNet(BaseModel):
@@ -313,14 +265,23 @@ class EulerNet(BaseModel):
                 "All interaction features must share the same embedding_dim in EulerNet."
             )
 
-        self.euler = EulerNetPaper(
-            embedding_dim=self.embedding_dim,
-            num_fields=self.num_fields,
-            num_layers=num_layers,
-            num_orders=num_orders,
-            use_implicit=use_implicit,
-            norm=norm,
+        self.num_layers = num_layers
+        self.num_orders = num_orders
+        self.mapping = ComplexSpaceMapping(self.embedding_dim, self.num_fields)
+        self.layers = nn.ModuleList(
+            [
+                EulerInteractionLayer(
+                    embedding_dim=self.embedding_dim,
+                    num_fields=(self.num_fields if i == 0 else self.num_orders),
+                    num_orders=self.num_orders,
+                    use_implicit=use_implicit,
+                    norm=norm,
+                )
+                for i in range(self.num_layers)
+            ]
         )
+        self.w = nn.Linear(self.num_orders * self.embedding_dim, 1, bias=False)
+        self.w_im = nn.Linear(self.num_orders * self.embedding_dim, 1, bias=False)
 
         if self.use_linear:
             if len(self.linear_features) == 0:
@@ -336,7 +297,7 @@ class EulerNet(BaseModel):
 
         self.prediction_layer = PredictionLayer(task_type=self.task)
 
-        modules = ["euler"]
+        modules = ["mapping", "layers", "w", "w_im"]
         if self.use_linear:
             modules.append("linear")
         self.register_regularization_weights(
@@ -354,7 +315,7 @@ class EulerNet(BaseModel):
         field_emb = self.embedding(
             x=x, features=self.interaction_features, squeeze_dim=False
         )
-        y_euler = self.euler(field_emb)
+        y_euler = self.euler_forward(field_emb)
 
         if self.use_linear and self.linear is not None:
             linear_input = self.embedding(
@@ -363,3 +324,11 @@ class EulerNet(BaseModel):
             y_euler = y_euler + self.linear(linear_input)
 
         return self.prediction_layer(y_euler)
+
+    def euler_forward(self, field_emb: torch.Tensor) -> torch.Tensor:
+        r, p = self.mapping(field_emb)
+        for layer in self.layers:
+            r, p = layer(r, p)
+        r_flat = r.reshape(r.size(0), self.num_orders * self.embedding_dim)
+        p_flat = p.reshape(p.size(0), self.num_orders * self.embedding_dim)
+        return self.w(r_flat) + self.w_im(p_flat)