Neural Networks & LLM Components · Hard

MMoE Task Routing

Mix shared expert outputs with an independent softmax gate for each task.

Task

Implement mmoe_routing(expert_outputs: list[list[list[float]]], task_gate_logits: list[list[list[float]]]) -> list[list[list[float]]]. Convert each task's gate logits into stable expert weights, then return a task-specific weighted mixture of the shared expert outputs for every example.

Requirements

expert_outputs has shape examples by experts by features.
task_gate_logits has shape tasks by examples by experts.
Every task uses its own gate distribution for every example.
Apply numerically stable softmax across experts only.
Use the same expert weights across all output features for one task-example pair.
Return shape tasks by examples by features.
Do not mutate either input.
Preserve an empty task list and preserve empty example lists for each task.

Example

expert_outputs = [
    [[1.0, 0.0], [0.0, 2.0]],
]
task_gate_logits = [
    [[0.0, 0.0]],
    [[0.0, 1.0]],
]

mmoe_routing(expert_outputs, task_gate_logits)
# [
#   [[0.5,        1.0       ]],
#   [[0.26894142, 1.46211716]],
# ]