tfm.nlp.layers.ExpertsChooseMaskedRouter

Masked matmul router using experts choose tokens assignment.

tfm.nlp.layers.ExpertsChooseMaskedRouter(
    num_experts: int,
    *,
    jitter_noise: float = 0.0,
    use_bias: bool = True,
    kernel_initializer: _InitializerType = _DEFAULT_KERNEL_INITIALIZER,
    bias_initializer: _InitializerType = _DEFAULT_BIAS_INITIALIZER,
    router_z_loss_weight: float = 0.0,
    export_metrics: bool = True,
    name: str = 'router',
    **kwargs
)

This router uses the same mechanism as in Mixture-of-Experts with Expert Choice (https://arxiv.org/abs/2202.09368): each expert selects its top expert_capacity tokens. An individual token may be processed by multiple experts or none at all.

Uses Keras add_loss() and add_metric() APIs.

num_experts Number of experts.
jitter_noise Amplitude of jitter noise applied to router logits.
use_bias Whether or not to use the bias term in computing the router weights.
kernel_initializer Kernel initializer for router weights.
bias_initializer Bias initializer for router weights.
router_z_loss_weight Weight for router_z_loss. Use non-zero values if running into training instability (esp. with dtype 'bfloat16' or lower).
export_metrics Whether to export metrics using Keras add_metric API.
name Layer name.
**kwargs Forwarded to super.

Methods

call

View source

call(
    inputs: tf.Tensor, *, expert_capacity: int, training: Optional[bool] = None
) -> RouterOutput

Computes dispatch and combine arrays for routing to experts.

Args
inputs Inputs to send to experts of shape [num_groups, tokens_per_group, hidden_dim].
expert_capacity Each group will send this many tokens to each expert.
training If true, apply jitter noise during routing. If not provided taken from tf.keras.backend.

Returns
Router indices or mask arrays (depending on router type).