r/LocalLLaMA · · 2 min read

How does MTP actually improve performance? Haven't understood the verification process

Mirrored from r/LocalLLaMA for archival readability. Support the source by reading on the original site.

Hi there, does anyone know how MTP works? I was trying to wrap my head
around it, but things aren't adding up.

For example, the input string I give it is "look up, " and I want it to output
"the sky is blue", except the MTP heads miss the mark and tell me "the sky is
green". How does the model verify that "green" is wrong in a decent amount of
time?

From what I understand, the first prefill step is standard and populates the
caches for the various layers using "look up, " as context. Then the last layer
does the courtesy of giving me N tokens instead of just one, using a lightweight model that attaches to the prediction head.

The next step is that I have "look up, the sky is green", which goes through
prefill. I reach the last layer (which takes care of the verification) and this
is where I don't quite get the trick.

The only way I would have to be 100% sure of the generated tokens would be to
autoregressively generate all N MTP tokens and compare them one by one. However,
this would make MTP slower than standard token generation, so they must be doing
something smarter than that.

Chatting with Gemini, it tells me "at the last layer you use the causal
attention mask, so you prefill one token at a time and verify that the output
embedding matches for each of the 4 tokens."

I really didn't understand this. Tokens are supposed to be generated one at a
time, so how is it possible that the prefill (which generally should take the
tokens I give it as "good" since they are the context preceding the
autoregressive generation) manages to verify that the embeddings I generated in
MTP are correct using only the information available at the last layer of the
model? By design of decoder-only LLMs, shouldn't it take each embedding and run
it from the beginning of the model to the end to get sufficiently informed
embeddings to be able to judge the work done by the MTP heads?

I hope I made myself clear; feel free to ask me questions if at any point you
spot flaws in my reasoning 😁

submitted by /u/Other-Astronaut-2868
[link] [comments]

Discussion (0)

Sign in to join the discussion. Free account, 30 seconds — email code or GitHub.

Sign in →

No comments yet. Sign in and be the first to say something.

More from r/LocalLLaMA