Content

1. Content
2. The Core Problem
3. A Fixed Model Versus An Adaptive Model
4. What Does Width Mean?
5. How Does The Gate Learn?
6. What Does Depth Mean?
7. Entropy: The Confidence Meter
8. Neural Grafting: Do Not Throw Away Weak Parts
9. The Full Training Story
10. What The Paper Reports
11. Why GPU Speedup Is Harder Than FLOP Reduction
12. How This Differs From Related Ideas
13. A Professor’s Mental Model
14. Final Takeaway

The Core Problem

Imagine a professor grading a stack of exam answers.

Some answers are obvious. One sentence is enough to know the student understood the idea. Some answers are ambiguous. The professor must read carefully, compare details, and think longer.

Now imagine a rule that says:

Spend exactly the same amount of time on every answer.

That is wasteful. Easy cases are over-processed. Hard cases may still need all the attention they can get.

Standard BERT behaves like that fixed-rule professor. Every input goes through the same full stack:

Input -> Layer 1 -> Layer 2 -> ... -> Layer 12 -> Prediction

The paper A Flexible BERT Model Enabling Width- and Depth-Dynamic Inference asks a better question:

Can BERT learn how much computation each input deserves?

The answer is yes. Flexible BERT makes computation adaptive along two axes:

1. Width: how much of each Transformer layer should be active?
2. Depth: how many layers should the input pass through before prediction?

The mental model is triage. Easy inputs get a quick checkup. Hard inputs get the full diagnostic path.

A Fixed Model Versus An Adaptive Model

Standard BERT gives every input the same computational budget. Flexible BERT gives every input a route.

That is the conceptual shift. The model is no longer just a fixed function. It becomes a small decision system wrapped around a Transformer.

What Does Width Mean?

A BERT layer is not one indivisible block. It contains attention heads and feed-forward network neurons. Width controls how many of those components are active.

Flexible BERT uses width ratios such as:

25%, 50%, 75%, 100%

At each layer, a small gating module looks at the current [CLS] representation and predicts which width ratio to use next.

A single input may take a route like this:

Layer 1  -> 50% width
Layer 2  -> 25% width
Layer 3  -> 75% width
Layer 4  -> 50% width
...

This matters because Flexible BERT does not choose one small model for all future inputs. It chooses width per input, per layer.

How Does The Gate Learn?

This is one of the most elegant parts of the paper.

The gate does not need a human to label examples as easy or hard. The model generates its own supervision by asking:

What is the smallest subnet that gets this example right?

For a training example, run the available width subnets:

If 25% already gets the answer right, the gate learns that the input is easy. If only 100% works, the gate learns that the input is hard.

This is a powerful teaching pattern: instead of manually defining input difficulty, let the model discover difficulty through which subnet succeeds.

What Does Depth Mean?

Depth is how many Transformer layers the input passes through.

Standard BERT-Base has 12 layers and always uses all of them. Flexible BERT adds intermediate classifiers called off-ramps. In the paper, off-ramps are placed after layer 6 and layer 9, with the final classifier after layer 12.

At each off-ramp, the model asks:

Am I confident enough to stop?

If yes, prediction happens early. If no, the input continues to deeper layers.

Entropy: The Confidence Meter

For classification tasks, the paper uses entropy as an uncertainty score.

Low entropy means the probability distribution is sharp:

Positive: 0.998
Negative: 0.002

That is a confident prediction. The input can exit early.

High entropy means the distribution is flat:

Positive: 0.52
Negative: 0.48

That is uncertain. The input should continue.

The rule is simple:

low entropy  -> confident -> exit
high entropy -> uncertain -> continue

For regression tasks such as STS-B, there is no softmax distribution. The paper instead uses prediction stability between layers: if the prediction is no longer changing much, the input can exit.

Neural Grafting: Do Not Throw Away Weak Parts

Flexible BERT is not only about routing at inference time. It also has to train many subnetworks inside one shared model.

That creates a problem. Small subnetworks use only a subset of attention heads and FFN neurons. Large subnetworks use more. Some components can become under-trained or inactive.

A pruning mindset would say:

If a component is weak, remove it.

Neural grafting says:

If a component is weak, rehabilitate it.

The model ranks attention heads and FFN neurons by importance. Then, inside a layer, weaker components receive a small weighted signal from stronger components.

This happens layer by layer. The same idea is applied to attention heads and FFN neurons.

The important intuition is that grafting is not replacement. The weak component keeps most of its identity and receives a small nudge from a stronger peer.

The Full Training Story

Flexible BERT is built in stages.

Under the hood, the subnet training uses knowledge distillation. The smaller paths learn from the behavior of larger or full teacher models.

What The Paper Reports

The headline result is not merely “smaller BERT is faster.”

The stronger claim is:

Adaptive BERT can improve the performance-cost tradeoff because it spends compute where it is useful.

Compared with BERT-Base, the paper reports:

• about 45% fewer computations
• about 2.3x CPU speedup
• +1.3 average GLUE points
• +2.2 SQuAD F1 points

The exact numbers depend on task and budget, but the direction is clear: dynamic routing gives a better Pareto frontier than simply picking one fixed compressed model.

Why GPU Speedup Is Harder Than FLOP Reduction

One subtle point in the paper is worth remembering.

Width reduction lowers FLOPs, but that does not always translate into proportional GPU latency gains. GPUs are good at dense parallel computation, and irregularly using parts of a layer can be harder to exploit.

Depth reduction is more directly useful for latency because skipping layers removes entire blocks of work.

So a practical serving lesson is:

width routing -> good for reducing arithmetic
depth routing -> often more visible for latency
width + depth -> best conceptual tradeoff

How This Differs From Related Ideas

Flexible BERT sits at the intersection of several model-efficiency ideas.

A Professor’s Mental Model

If I were teaching this in a classroom, I would ask students to remember one image:

Flexible BERT is a hospital triage desk attached to BERT.

The patient is the input.

The gate asks:

How much capacity does this patient need at this stage?

The off-ramp asks:

Is the diagnosis confident enough to stop?

Neural grafting asks:

Are some departments under-trained, and can stronger departments help them recover?

This gives us the three central mechanisms:

1. Width routing decides how much of a layer to use.
2. Depth routing decides when to stop.
3. Neural grafting keeps the available subnetworks healthy.

Final Takeaway

Flexible BERT changes the compression question.

Instead of asking:

How small can we make BERT?

it asks:

How much BERT does this input need?

That is the deeper idea. Future efficient models should not only be smaller. They should be adaptive. They should learn when to think more, and just as importantly, when not to.