mixture_of_experts

mixture_of_experts

Mixture-of-experts loss wrapper with routing regularization.

MoELoss

MoELoss(base_loss, entropy_weight=0.0, load_balance_weight=0.0)

Bases: BaseLoss

Wrapper for any BaseLoss to add MoE expert usage metric from routing weights.

Initialize MoE wrapper.

Parameters:

Name	Type	Description	Default
base_loss	BaseLoss	Any BaseLoss instance to wrap (e.g., RegressionLoss(...))	required
entropy_weight	float	Weight for per-example routing entropy. Penalizes peaky-per-example routing. Pushes each example's routing distribution toward uniform, which prevents one example from being routed to a single expert with probability 1.	0.0
load_balance_weight	float	Weight for Switch-Transformer-style load-balancing term. Penalizes batch-level imbalance in expert usage. The term is K * sum_k f_k * P_k where f_k is the fraction of examples whose argmax routes to expert k and P_k is the mean routing weight for expert k across the batch. Minimum value 1.0 is reached when usage is uniform across the batch. Crucially, this allows per-example routing to be peaky (so experts can specialize) while still forcing every expert to be used by some examples (so no expert dies). Use this when entropy alone produces dead experts.	0.0

Source code in src/versatil/metrics/losses/mixture_of_experts.py

def __init__(
    self,
    base_loss: BaseLoss,
    entropy_weight: float = 0.0,
    load_balance_weight: float = 0.0,
):
    """Initialize MoE wrapper.

    Args:
        base_loss: Any BaseLoss instance to wrap (e.g., RegressionLoss(...))
        entropy_weight: Weight for per-example routing entropy.
            Penalizes peaky-per-example routing. Pushes each example's routing
            distribution toward uniform, which prevents one example from being
            routed to a single expert with probability 1.
        load_balance_weight: Weight for Switch-Transformer-style load-balancing
            term. Penalizes batch-level imbalance in expert usage. The term is
            ``K * sum_k f_k * P_k`` where ``f_k`` is the fraction of examples
            whose argmax routes to expert k and ``P_k`` is the mean routing
            weight for expert k across the batch. Minimum value 1.0 is reached
            when usage is uniform across the batch. Crucially, this allows
            per-example routing to be peaky (so experts can specialize) while
            still forcing every expert to be used by some examples (so no
            expert dies). Use this when entropy alone produces dead experts.
    """
    super().__init__()
    self.base_loss = base_loss
    self.entropy_weight = entropy_weight
    self.load_balance_weight = load_balance_weight

weights property

weights

Getter that returns dictionary with weight keys and scalar coefficients, plus the wrapped base_loss weight structure nested under base_loss.

set_weights

set_weights(new_weights)

Setter that updates the weight scalar coefficients and delegates base_loss to the wrapped loss.

Source code in src/versatil/metrics/losses/mixture_of_experts.py

def set_weights(self, new_weights: WeightsDictionary) -> None:
    """Setter that updates the weight scalar coefficients and delegates
    ``base_loss`` to the wrapped loss."""
    self._validate_weights(new_weights)
    self.entropy_weight = new_weights["entropy_weight"]
    self.load_balance_weight = new_weights["load_balance_weight"]
    self.base_loss.set_weights(new_weights["base_loss"])

get_callbacks

get_callbacks(experiment_config)

Provide expert usage monitoring callback.

Source code in src/versatil/metrics/losses/mixture_of_experts.py

def get_callbacks(self, experiment_config: ExperimentConfig) -> list:
    """Provide expert usage monitoring callback."""
    return [ExpertUsageCallback(log_every_n_epochs=1)]

get_required_keys

get_required_keys()

Union of base loss keys plus routing weight.

Source code in src/versatil/metrics/losses/mixture_of_experts.py

def get_required_keys(self) -> set[str]:
    """Union of base loss keys plus routing weight."""
    return self.base_loss.get_required_keys() | {
        DecoderOutputKey.ROUTING_WEIGHTS.value
    }

forward

forward(predictions, targets, is_pad=None)

Passthrough base loss, then add expert_usage and optional entropy/load-balance terms.

Source code in src/versatil/metrics/losses/mixture_of_experts.py

def forward(
    self,
    predictions: dict[str, torch.Tensor],
    targets: dict[str, torch.Tensor],
    is_pad: torch.Tensor | None = None,
) -> LossOutput:
    """Passthrough base loss, then add expert_usage and optional entropy/load-balance terms."""
    predictions = self._add_weighted_mean_predictions(predictions)
    base_output: LossOutput = self.base_loss(predictions, targets, is_pad)
    metadata = base_output.metadata if base_output.metadata is not None else {}
    component_losses = dict(base_output.component_losses)
    pi = predictions[DecoderOutputKey.ROUTING_WEIGHTS.value]  # (B, K) or (B, T, K)
    total_loss = base_output.total_loss
    if self.entropy_weight != 0.0:
        entropy = -(pi * torch.log(pi + 1e-8)).sum(dim=-1)  # (B,) or (B, T)
        if entropy.dim() == 2:
            entropy_mean = reduce_loss_with_padding(
                entropy, is_pad, reduction="mean"
            )
        else:
            entropy_mean = entropy.mean()
        component_losses[f"{MetricKey.EXPERTS_ENTROPY.value}"] = entropy_mean
        total_loss = total_loss - self.entropy_weight * entropy_mean
    if self.load_balance_weight != 0.0:
        load_balance = self._compute_load_balance(pi=pi, is_pad=is_pad)
        component_losses[f"{MetricKey.EXPERTS_LOAD_BALANCE.value}"] = load_balance
        total_loss = total_loss + self.load_balance_weight * load_balance
    expert_usage = pi.mean(
        dim=list(range(pi.ndim - 1))
    )  # Mean over all but last dim, which is num_experts
    metadata[MetadataKey.EXPERT_USAGE.value] = expert_usage
    return LossOutput(
        total_loss=total_loss,
        component_losses=component_losses,
        metadata=metadata,
    )