LatentMoE: must experts be as wide as the backbone?

Sending a token to eight experts means moving its representation to those experts. Even small expert computations can spend substantial time reading weights and moving data. What if the experts could work with a narrower representation?

That is the idea behind LatentMoE. The routed experts' input and output become narrower; the backbone does not. Start with the paths and budget table for the main idea. The later sections cover stability and balancing.

Which part becomes narrower?

One important connection is easy to omit: the router can select experts from the original width-128 input, not the width-64 latent. The vector being dispatched and the input used to score experts are different objects. NVIDIA's original explanation makes this distinction explicit.

Using column vectors, let $D$ project down, $U$ project up, and $\mathcal T(x)$ denote the chosen experts. A basic version is

$$ \begin{gathered} z=Dx,\\ u=\sum_{i\in\mathcal T(x)}g_i(x)E_i(z),\\ y=S(x)+Uu. \end{gathered} $$

$S$ combines the shared experts; the outer residual is omitted. Backbone width $d$, latent width $\ell$, and expert intermediate width $f$ are different dimensions. An expert maps $\ell\rightarrow f\rightarrow\ell$. Not every dimension has been halved.

Does the same latent vector imply the same output?

Not necessarily: the router still sees the original input. In a two-dimensional toy example, keep only the first coordinate as the latent. Use the original coordinates as routing scores, select the larger score, and normalize the selected expert's weight to 1. Expert A doubles its input; expert B negates it. The up-projection copies the scalar output into both coordinates. Leave out shared experts and the residual for now.

Both inputs send the same scalar 1 to the experts, but different routing decisions produce different outputs. Feeding the compressed representation to the router instead would erase this distinction.

def routed_example(vector):
    latent = vector[0]
    selected = 0 if vector[0] >= vector[1] else 1
    combined = (2 * latent, -latent)[selected]
    return latent, selected, (combined, combined)

assert routed_example((1, 0)) == (1, 0, (2, 2))
assert routed_example((1, 2)) == (1, 1, (-1, -1))

Three different ways to make something smaller

LatentMoE and MLA can coexist, but LatentMoE does not itself compress attention's KV cache. For why MLA separates its RoPE branch, see matrix absorption and position.

Does the saving pay for more experts?

Work through the savings and added costs with small dimensions: $d=128,f=256$, eight SwiGLU experts, and two active experts per token. Ignoring bias, each expert has $3df$ weights. LatentMoE changes this to $3\ell f$ and adds $2d\ell$ projection weights.

The third row matches the first row's expert-only budget, but the projections still cost something. Parameters rise about 2.1%; main MACs rise about 8.3%. The change is not free. Routing, shared experts, normalization, activations, and attention are excluded. Counting multiply and add separately, one MAC is about two FLOPs.

def latent_budget(width, latent, intermediate, experts, top_k, tokens, item_bytes):
    values = (width, latent, intermediate, experts, top_k, tokens, item_bytes)
    if any(type(value) is not int or value < 1 for value in values):
        raise ValueError("Expected positive integer dimensions and counts")
    if latent > width or top_k > experts:
        raise ValueError("Invalid latent width or expert selection")
    projections = 0 if latent == width else 2 * width * latent
    per_expert = 3 * latent * intermediate
    return {
        "parameters": experts * per_expert + projections,
        "mac_per_token": top_k * per_expert + projections,
        "slot_bytes": 2 * tokens * top_k * latent * item_bytes,
    }

base = latent_budget(128, 128, 256, 8, 2, 32, 2)
small = latent_budget(128, 64, 256, 8, 2, 32, 2)
more = latent_budget(128, 64, 256, 16, 4, 32, 2)
assert [row["parameters"] for row in (base, small, more)] == [786432, 409600, 802816]
assert [row["mac_per_token"] for row in (base, small, more)] == [196608, 114688, 212992]
assert [row["slot_bytes"] for row in (base, small, more)] == [32768, 16384, 32768]

Halving the width but doubling top-k brings the slot bytes back to the baseline. The question becomes whether quality improves at that cost, not just how much a dimension shrank. The original LatentMoE paper distinguishes compute from bandwidth and communication constraints. The numbers above are independently constructed examples, not its benchmarks.

Kimi K3: what else needs attention?

K3 selects 16 of 896 routed experts per token. Its official release calls the design Stable LatentMoE. Report §2.3 specifies a width-7168 backbone, width-3584 routed path, and two full-width shared experts; the router scores the full input. Dispatching a 3584-dimensional vector does not mean the router must read one.

It adds RMSNorm between the combined expert output and up-projection, uses SiTU-GLU to control internal activation range, and uses Quantile Balancing for load. These address output scale, internal products, and uneven dispatch, respectively.

Why the normalization cannot simply move

Consider two expert outputs, [2, 0] and [0, 1], weighted equally. Ignore epsilon and learned gains:

The first preserves the 2:1 direction; the second changes it to 1:1. RMSNorm is not a linear projection that can move through a sum. An implementation can have the right shapes and still compute a different model.

SiTU has a related intuition: multiplying two large coordinates makes outliers worse, so smoothly cap both branches before multiplying. This is neither a constant replacement for every activation nor a guarantee against all network overflow. Learned projections can still amplify outputs, and gradients can become small in saturated regions.

Quantile Balancing: measure now, use next step

Six tokens, three experts, and top-1 produce six assignments, giving an average target of two per expert. Average load alone does not tell us how to adjust a bias. We also need to know how far each expert is from being selected.

K3 uses the current biased $(k+1)$-th score as a cutoff and expert-wise score margins to estimate the next bias. Mixture weights still use raw affinities. The bias takes effect next step and is frozen for inference. Global quantiles are approximated with histograms. Report §2.3.3

There are two important limits. A quantile calculated against old cutoffs does not guarantee perfectly balanced top-k assignments after every bias changes; the next batch can differ too. And “aux-loss-free” does not mean ordinary task gradients update the bias. A separate state-update rule still exists.

def fixed_cutoff_bias(margins, target):
    import math

    if not margins or any(not math.isfinite(value) for value in margins):
        raise ValueError("Expected finite margins")
    if type(target) is not int or not 0 < target < len(margins):
        raise ValueError("Expected an interior integer target")
    ordered = sorted(margins, reverse=True)
    return -ordered[target]

margins = [0.31, 0.16, 0.07, -0.08, -0.19, -0.34]
bias = fixed_cutoff_bias(margins, 2)
assert sum(margin + bias > 0 for margin in margins) == 2

How would you decide whether to use it?

First identify the bottleneck. Small-batch decode may mainly read weights. Larger distributed batches may wait for all-to-all or overloaded experts. Top-k and parameter count alone cannot distinguish these cases.

Nemotron 3 Super also uses LatentMoE, with different dimensions: its report specifies a $4096\rightarrow1024$ routed path. Nemotron 3 Super §2.1.1 Compression ratio is a design choice, not part of the definition. Published gains belong to those experiments. This page checks arithmetic and mathematical relationships, not model training or multi-GPU performance.

If you can explain the unchanged backbone, the difference between router input and dispatch payload, and the cost of extra projections, you have the main idea. Continue with system costs to turn these dimensions into storage and time on actual hardware.