DEV Community

Cover image for Google’s Gemini Handles a Million Tokens. Here’s the Architecture Trick That Makes It Possible
Ankita Maji
Ankita Maji

Posted on Originally published at Medium

Google’s Gemini Handles a Million Tokens. Here’s the Architecture Trick That Makes It Possible

Most AI explainers skip Mixture of Experts because it sounds complicated. It isn’t. And understanding it explains why frontier AI costs what it does.

Scaling a neural network used to mean one thing: make every layer bigger, train on more data, pay more for compute. That worked — until it stopped working economically.

Training a model where every single parameter activates for every single input token is expensive in a way that scales brutally. Double the parameters, roughly double the compute cost per token. At a hundred billion parameters, that math becomes a problem that no amount of funding fully solves.

Mixture of Experts — MoE — is the architectural answer to that problem. It’s the reason Gemini can handle a million-token context without the inference cost becoming absurd. It’s the reason Mixtral 8x7B punches far above its weight class. And it’s one of the most important ideas in modern AI that most explainers skip over because the name sounds intimidating.

It shouldn’t. The core idea takes about two minutes to understand.

The dense model problem

Every transformer layer in a standard — “dense” — model works the same way. A token comes in. It passes through every neuron in the feed-forward layer. Every parameter participates. Every computation happens.

This is fine when models are small. It becomes a serious problem at scale.

Consider: GPT-3 has 175 billion parameters. For every single token it processes — every word in your prompt, every word it generates — all 175 billion parameters are active. The compute cost per token is proportional to the total parameter count. You cannot separate “which parameters matter for this token” from “all parameters.”

This is the dense model tax. You pay for every parameter on every token, regardless of whether those parameters are useful for what you’re currently processing.

The insight: not every expert needs to weigh in on everything

Here’s the Mixture of Experts idea, stated plainly.

Instead of one large feed-forward layer where every neuron activates for every token, replace it with multiple smaller feed-forward networks — call them experts. Add a lightweight routing network — the gating network — that looks at each token and decides which experts should handle it.

For each token, only a small number of experts activate. The rest sit idle. The token gets processed by its assigned experts and moves on.

The result: a model can have a very large total parameter count — all the experts combined — while only activating a fraction of those parameters for any given token. Total parameters and active parameters are now two different numbers, and you get to optimise them independently.

Mixtral 8x7B has 8 experts per layer, each with 7 billion parameters — 56 billion total. But for each token, only 2 experts activate. The effective compute per token is closer to a 14B dense model, not a 56B one. You get the knowledge capacity of a larger model at the inference cost of a smaller one.


How the routing actually works

The gating network is the part most explainers skip. It matters.

Become a Medium member
For each token, the gating network produces a score for every expert — a number representing how relevant that expert is for this particular token. The top-K experts by score are selected. Their outputs are weighted by their scores and combined.

In practice, K is usually 1 or 2. Top-1 routing means each token goes to exactly one expert. Top-2 means two experts process it and their outputs are blended.

The gating network is tiny — just a linear layer — and it learns during training which kinds of tokens different experts should handle. Nobody hand-labels which expert handles which topic. The routing emerges from training. Some experts end up specialising in syntax, others in factual recall, others in code. The model discovers this structure on its own.

One important practical problem: load balancing. If the gating network consistently routes most tokens to the same two experts, those experts get overloaded while the rest sit idle. This defeats the purpose. Training MoE models requires an auxiliary loss function that penalises unbalanced routing — pushing the model to distribute tokens more evenly across experts. Getting this right is one of the main engineering challenges in MoE training.


Why this matters for long-context models like Gemini

Gemini’s 1-million-token context window is not just a transformer scaling achievement — it requires careful management of compute cost per token. If every token in a million-token context activated a full dense model, the compute per forward pass would be staggering.

MoE makes long-context processing economically viable by keeping active parameter count per token manageable even as total model capacity grows. The million-token context is processed token by token through sparse experts — not through a wall of fully-active dense computation.

This is also why MoE models have a different cost structure than dense models. Training cost is higher — you’re training all the experts, which requires more GPU memory and more careful engineering. But inference cost per token is lower, because only a fraction of parameters activate. For a product like Gemini that serves millions of queries per day, the inference cost saving dominates.


The engineering challenges nobody talks about

MoE is not a free lunch. Three real challenges come with it.

Memory. All experts must be loaded into memory even though only a fraction activate per token. A model with 56B total parameters requires memory for all 56B, even though only 14B activate at once. For deployment on devices with limited memory — exactly the kind of edge nodes MOSAIC targets — this is a significant constraint.

Communication overhead. In distributed training and inference, different experts often live on different GPUs. Routing tokens to the right expert means sending data across GPU interconnects. This all-to-all communication pattern is expensive and becomes a bottleneck at scale. It’s one reason MoE models are harder to serve efficiently than dense models of equivalent active parameter count.

Training instability. MoE models are notoriously harder to train than dense models. The routing network and the experts need to co-adapt — and early in training, before the routing has stabilised, the loss can spike erratically. Getting MoE training to converge reliably at scale is a significant engineering challenge that Google, Mistral, and others have invested heavily in solving.


What this means for the future

Mixture of Experts is not a temporary hack. It’s a fundamental architectural insight that the field has converged on for scaling frontier models.

The reason is straightforward: the alternative — scaling dense models indefinitely — has diminishing returns and growing costs. MoE offers a way to increase model capacity without proportionally increasing inference cost. For companies serving AI at scale, that trade-off is compelling enough that every major frontier model is now either using MoE or actively exploring it.

Understanding MoE means understanding why AI products are priced the way they are, why some models feel faster than their parameter count suggests they should be, and why the gap between “total parameters” and “active parameters” is one of the most important numbers in AI infrastructure.

Next time someone quotes a model’s parameter count as a measure of its capability — remember that number tells you less than you think.

Drop a comment or connect on LinkedIn — always open to feedback and pushback.

Top comments (0)