Mixture of Experts

network/body/embedding.py

@@ -263,7 +263,11 @@ class PETBody(nn.Module):
     def __init__(
             self, num_feat, num_keep, feature_drop, projection_dim, local, K, num_local,
             num_layers, num_heads, drop_probability, talking_head, layer_scale,
-            layer_scale_init, dropout, mode, use_adapter: bool = False
+            layer_scale_init, dropout, mode, use_adapter: bool = False,
+            use_moe: bool = False, moe_base_num_experts: int = 4,
+            moe_base_select_top_k: int = 2, moe_num_shared_experts: int = 0,
+            moe_expert_segmentation_factor: int = 1, moe_scale_expert_dim: bool = False,
+            moe_alpha: float = 0.01, moe_cz: float = 0.0, moe_use_router_noise: bool = False
     ):
         super().__init__()
         self.num_keep = num_keep
@@ -295,7 +299,12 @@ def __init__(
         self.transformer_blocks = nn.ModuleList([
             TransformerBlockModule(
                 projection_dim, num_heads, dropout, talking_head, layer_scale, layer_scale_init,
-                drop_probability
+                drop_probability, use_moe=use_moe, moe_base_num_experts=moe_base_num_experts,
+                moe_base_select_top_k=moe_base_select_top_k, moe_num_shared_experts=moe_num_shared_experts,
+                moe_expert_segmentation_factor=moe_expert_segmentation_factor,
+                moe_scale_expert_dim=moe_scale_expert_dim,
+                moe_alpha=moe_alpha, moe_cz=moe_cz,
+                moe_use_router_noise=moe_use_router_noise
             )
             for _ in range(num_layers)
         ])
@@ -350,16 +359,21 @@ def forward(self,
             encoded = local_features + encoded  # Combine with original features
 
         skip_connection = encoded
+        moe_l_aux = encoded.new_zeros(())
+        moe_cz_lz = encoded.new_zeros(())
         for itransformer, transformer_block in enumerate(self.transformer_blocks):
             encoded = transformer_block(
                 x=encoded,
                 mask=mask,
                 attn_mask=attn_mask
             )
+            moe_l_aux += transformer_block.moe_l_aux.to(encoded.device)
+            moe_cz_lz += transformer_block.moe_cz_lz.to(encoded.device)
             if self.use_adapter:
                 encoded = self.adapters[itransformer](encoded)
                 encoded = encoded * mask.float()
-
+        self.moe_l_aux = moe_l_aux
+        self.moe_cz_lz = moe_cz_lz
 
         return torch.add(encoded, skip_connection)
 

network/evenet_model.py

@@ -21,6 +21,7 @@
 from torch import Tensor, nn
 from typing import Dict, Optional, Any, Union
 import re
+import logging
 
 
 class EveNetModel(nn.Module):
@@ -155,6 +156,15 @@ def __init__(
             layer_scale_init=pet_config.layer_scale_init,
             dropout=pet_config.dropout,
             mode=pet_config.mode,
+            use_moe=pet_config.use_moe,
+            moe_base_num_experts=pet_config.moe_base_num_experts,
+            moe_base_select_top_k=pet_config.moe_base_select_top_k,
+            moe_num_shared_experts=pet_config.moe_num_shared_experts,
+            moe_expert_segmentation_factor=pet_config.moe_expert_segmentation_factor,
+            moe_scale_expert_dim=pet_config.moe_scale_expert_dim,
+            moe_alpha=pet_config.moe_alpha,
+            moe_cz=pet_config.moe_cz,
+            moe_use_router_noise=pet_config.moe_use_router_noise,
         )
 
         # [2] Classification + Regression + Assignment Body
@@ -317,6 +327,41 @@ def __init__(
             ("deterministic", self.include_classification or self.include_assignment or self.include_regression or self.include_segmentation),
         ]
 
+        self._log_backbone_setup()
+
+    def _log_backbone_setup(self) -> None:
+        logger = logging.getLogger(__name__)
+        pet_cfg = self.network_cfg.Body.PET
+
+        pretrain_path = getattr(getattr(self.options, "Training", None), "pretrain_model_load_path", None)
+        if pretrain_path:
+            logger.info(f"[Backbone] Pretrain path       : {pretrain_path}")
+        else:
+            logger.warning("[Backbone] No pretrain_model_load_path set — training from scratch.")
+
+        logger.info(
+            f"[Backbone] PET config          : "
+            f"layers={pet_cfg.num_layers}, "
+            f"heads={pet_cfg.num_heads}, "
+            f"dim={pet_cfg.hidden_dim}, "
+            f"mode={pet_cfg.mode}"
+        )
+        logger.info(
+            f"[Backbone] MoE enabled         : {pet_cfg.use_moe}"
+        )
+        if pet_cfg.use_moe:
+            logger.info(
+                f"[Backbone] MoE config          : "
+                f"num_experts={pet_cfg.moe_base_num_experts}, "
+                f"top_k={pet_cfg.moe_base_select_top_k}, "
+                f"shared_experts={pet_cfg.moe_num_shared_experts}, "
+                f"seg_factor={pet_cfg.moe_expert_segmentation_factor}, "
+                f"scale_dim={pet_cfg.moe_scale_expert_dim}, "
+                f"alpha={pet_cfg.moe_alpha}, "
+                f"cz={pet_cfg.moe_cz}, "
+                f"router_noise={pet_cfg.moe_use_router_noise}"
+            )
+
     def forward(
             self, x: Dict[str, Tensor], time: Tensor,
             progressive_params: dict = None,
@@ -461,6 +506,8 @@ def forward(
 
         full_input_point_cloud = None
         full_global_conditions = None
+        moe_l_aux_total = torch.zeros((), device=input_point_cloud.device, dtype=input_point_cloud.dtype)
+        moe_cz_lz_total = torch.zeros((), device=input_point_cloud.device, dtype=input_point_cloud.dtype)
 
         for schedule_name, flag  in schedules:
             if not flag:
@@ -537,6 +584,8 @@ def forward(
                 time=full_time,
                 time_masking=time_masking
             )
+            moe_l_aux_total += self.PET.moe_l_aux
+            moe_cz_lz_total += self.PET.moe_cz_lz
 
             if schedule_name == "deterministic" or schedule_name == "generation":
                 ######################################
@@ -655,7 +704,9 @@ def forward(
             # "full_global_conditions": full_global_conditions,
             "alpha": alpha,
             "segmentation-mask": outputs.get("deterministic", {}).get("segmentation-out", {}).get("pred_masks", None),
-            "segmentation-aux": outputs.get("deterministic", {}).get("segmentation-out", {}).get("aux_outputs", None)
+            "segmentation-aux": outputs.get("deterministic", {}).get("segmentation-out", {}).get("aux_outputs", None),
+            "L_aux": moe_l_aux_total,
+            "cz_Lz": moe_cz_lz_total,
         }
 
     def predict_diffusion_vector(