mixture_of_experts

mixture_of_experts

Base Mixture of Experts module with shared gating and routing logic.

BaseMixtureOfExperts

BaseMixtureOfExperts(num_experts, device, gating_input_dim=None, gating_activation_function=value, gating_hidden_dims=None, routing_type=value, top_k=2, temperature=1.0, learnable_temperature=False, gating_dropout=0.1, gating_normalization=True)

Bases: Module

Base class for Mixture of Experts with shared gating and routing logic.

Handles: - Gating network creation and management - Temperature-scaled routing weight computation - Expert specialization analysis - Routing strategies (soft, top-k)

Initialize Mixture of Experts routing logic.

Parameters:

Name	Type	Description	Default
num_experts	int	Number of expert models	required
device	str	Device to place gating network on	required
gating_input_dim	int | None	Input dimension for gating network (None for external routing)	None
gating_activation_function	str	Activation function for gating network	value
gating_hidden_dims	list[int] | None	Hidden layer dimensions for gating network	None
routing_type	str	Routing strategy ("soft" or "top_k")	value
top_k	int	Number of experts to use for top-k routing	2
temperature	float	Temperature for softmax scaling of routing weights	1.0
learnable_temperature	bool	Whether temperature should be a learnable parameter	False
gating_dropout	float	Dropout rate in gating network	0.1
gating_normalization	bool	Whether to normalize inputs to gating network	True

Source code in src/versatil/models/decoding/mixture_of_experts.py

def __init__(
    self,
    num_experts: int,
    device: str,
    gating_input_dim: int | None = None,
    gating_activation_function: str = ActivationFunction.RELU.value,
    gating_hidden_dims: list[int] | None = None,
    routing_type: str = MoERoutingType.SOFT.value,
    top_k: int = 2,
    temperature: float = 1.0,
    learnable_temperature: bool = False,
    gating_dropout: float = 0.1,
    gating_normalization: bool = True,
):
    """Initialize Mixture of Experts routing logic.

    Args:
        num_experts: Number of expert models
        device: Device to place gating network on
        gating_input_dim: Input dimension for gating network (None for external routing)
        gating_activation_function: Activation function for gating network
        gating_hidden_dims: Hidden layer dimensions for gating network
        routing_type: Routing strategy ("soft" or "top_k")
        top_k: Number of experts to use for top-k routing
        temperature: Temperature for softmax scaling of routing weights
        learnable_temperature: Whether temperature should be a learnable parameter
        gating_dropout: Dropout rate in gating network
        gating_normalization: Whether to normalize inputs to gating network
    """
    nn.Module.__init__(self)
    if num_experts == 0:
        raise ValueError("Must provide at least one expert")
    valid_routing_types = [e.value for e in MoERoutingType]
    if routing_type not in valid_routing_types:
        raise ValueError(
            f"Invalid routing_type: {routing_type}. Expected one of {valid_routing_types}"
        )

    if num_experts < 1:
        raise ValueError(f"num_experts must be positive, got {num_experts}.")
    if not 1 <= top_k <= num_experts:
        raise ValueError(
            f"top_k must be in [1, num_experts={num_experts}], got {top_k}."
        )
    if temperature <= 0.0:
        raise ValueError(f"temperature must be positive, got {temperature}.")
    self.num_experts = num_experts
    self.routing_type = routing_type
    self.top_k = top_k
    self.has_gating_network = gating_input_dim is not None
    self.gating_network: nn.Sequential | None
    if self.has_gating_network:
        self.gating_network = self._build_gating_network(
            input_dimension=gating_input_dim,
            hidden_dimensions=gating_hidden_dims,
            activation=gating_activation_function,
            dropout=gating_dropout,
            normalization=gating_normalization,
            device=device,
        )
    else:
        self.gating_network = None
    if learnable_temperature:
        self.temperature = nn.Parameter(
            torch.tensor(temperature, dtype=torch.float32), requires_grad=True
        )
    else:
        self.register_buffer(
            "temperature", torch.tensor(temperature, dtype=torch.float32)
        )

compute_routing_weights

compute_routing_weights(features)

Compute routing weights from input or external source.

Parameters:

Name	Type	Description	Default
features	Tensor	Input tensor for gating network	required

Returns:

Type	Description
Tensor	Normalized routing weights (B, [horizon,] num_experts)

Source code in src/versatil/models/decoding/mixture_of_experts.py

def compute_routing_weights(self, features: torch.Tensor) -> torch.Tensor:
    """Compute routing weights from input or external source.

    Args:
        features: Input tensor for gating network

    Returns:
        Normalized routing weights (B, [horizon,] num_experts)

    """
    logits = self.gating_network(features) if self.has_gating_network else features
    logits = logits / self.temperature
    return F.softmax(logits, dim=-1)

get_expert_specialization

get_expert_specialization(gating_feature)

Analyze expert usage patterns.

Parameters:

Name	Type	Description	Default
gating_feature	Tensor	Input gating feature	required

Returns:

Type	Description
dict[str, Tensor]	Dictionary with expert_usage, routing_entropy, top_expert_confidence

Note

This method uses torch.no_grad() for efficiency but does not modify the model's training state. Caller should set model to eval mode if needed.

Source code in src/versatil/models/decoding/mixture_of_experts.py

def get_expert_specialization(
    self,
    gating_feature: torch.Tensor,
) -> dict[str, torch.Tensor]:
    """Analyze expert usage patterns.

    Args:
        gating_feature: Input gating feature

    Returns:
        Dictionary with expert_usage, routing_entropy, top_expert_confidence

    Note:
        This method uses torch.no_grad() for efficiency but does not modify
        the model's training state. Caller should set model to eval mode if needed.
    """
    with torch.no_grad():
        weights = self.compute_routing_weights(gating_feature)
        expert_usage = weights.mean(dim=tuple(range(weights.ndim - 1)))
        entropy = -(weights * torch.log(weights + 1e-8)).sum(dim=-1).mean()
        top_expert_confidence = weights.max(dim=-1)[0].mean()
    return {
        DecoderOutputKey.EXPERT_USAGE.value: expert_usage,
        DecoderOutputKey.ROUTING_ENTROPY.value: entropy,
        DecoderOutputKey.TOP_EXPERT_CONFIDENCE.value: top_expert_confidence,
    }