Why does LLM inference cost scale in both input tokens and output tokens?

Question

EDIT

This question was flawed. See my answer with help from commenters.

---

Original question

This question has been asked in other forums [1] [2] but I'm not sure I understand the claims, which are (EDIT: the following are based on my faulty assumption that pad tokens are added up to the maximum context window):

1. Each forward pass takes less resources when more of the context window is padding.
2. Forward passes are run on the input tokens.
3. Forward passes with fewer non-pad input tokens are smaller tensor operations.

Hypothesis 1 seems the most plausible to me from a performance engineering standpoint (sparse math, etc). Does it fall out naturally from just writing basic JAX code or would it require manual optimization (if so, what tricks can be used?)? There does seem to be some research on this.

Hypothesis 2 and 3 seem wrong based on my surface-level understanding of the Transformer architecture.

I've tried a few open-source LLMs locally and neither on those nor ChatGPT have I noticed any difference in latency based on how much text was in the context window. But I haven't done actual rigorous benchmarking yet.

The reason this is relevant is due to document lookup-based applications. Looking at the OpenAI cookbook for Q&A using embeddings, I can see that:

• most of the token usage seems to come from the long Wikipedia article pasted into the context
• the price per query is quite prohibitive to frequent usage:

For gpt-3.5-turbo using ~1,000 tokens per query, it costs ~0.002 per query, or ~500 queries per dollar (as of Apr 2023) For gpt-4, again assuming ~1,000 tokens per query, it costs ~0.03 per query, or ~30 queries per dollar (as of Apr 2023)

My prior intuition would have been that optimal usage of LLMs would be to keep the context filled with inexpensive text (e.g. from NN search and/or cheaper LMs) and to have the LLM generate terse responses. But the input token cost model changes the strategy, as it means that users need to be sparing about the size and quantity of documents that they paste into the context window.

Answer 1

Transformers can be run with a variable dimension N_INPUT as pointed out by @pcpthm's comment. That means that the answer to "Why do transformers have a fixed input length?" are wrong, and the answers to "How do transformers handle arbitrary-length input?" are more correct.

If the input of dimension N_INPUT * N_VOCAB and N_INPUT is variable then that simply means Q, K, V, Z have one variable dimension (N_INPUT * D_ATTN_HEAD) and input/output of MLP are just N_INPUT * WIDTH.

For some reason I thought the hidden layers would be fixed size and/or require a fixed size input, but this is not the case since it turns out that N_INPUT is always the first dimension, except in the result of Q * K^T which has dimension N_INPUT * N_INPUT, but is only multiplied on the right by V which has first dimension N_INPUT.

In the NNs that I'd learned about previously, N_INPUT is the second dimension so even if one were to put in a variable input, all of the hidden layers would still be fixed-size.

Particularly what I did not realize is that the "context window" is not actually a part of the model (except in the positional encoding). It seems to be purely a training time / inference time restriction.