Inside the Final Layer: Logits, Sampling, and Structured Outputs in LLMs
How raw model scores become words, valid JSON, and filtered responses one token at a time.Most people treat LLM output as a black box. You send a prompt, text comes back, and whatever happens in between feels like magic.It is not magic. It is a short, well-defined chain of steps: hidden states, a…
How raw model scores become words, valid JSON, and filtered responses one token at a time.Most people treat LLM output as a black box. You send a prompt, text comes back, and whatever happens in between feels like magic.It is not magic. It is a short, well-defined chain of steps: hidden states, a linear layer, softmax, sampling, and then a set of engineering choices layered on top like temperature, top-p, and content filters.Once you can see that chain clearly, a lot of things stop being mysterious. Why the same prompt gives different answers each time. Why “temperature” changes how creative a model feels. Why a model can guarantee valid JSON, and how that guarantee actually works. Why safety filtering is not one thing but three separate layers doing different jobs.This article walks through that full pipeline start to finish, in the order it actually happens. I am talking about both the Generation layer and the Serving layer. The generation layer decides HOW to pick a token. The serving layer decides WHICH tokens are even allowed and WHETHER the response goes out.Generation layer — Temperature, Top-p/Top-K, Sampling strategyServing Layer — Json Schema(Logits masking),Content FilteringThe Final Transformer Block- What Comes OutBy the time your prompt reaches the last transformer block, it has already been turned into a sequence of hidden states, one vector per token. Each hidden state is the model’s internal summary of that token, given everything that came before it.The shape of this output is[seq_len, d_model], one row per token in your sequence, each row a vector of length d_model (4096, 8192, whatever your model's hidden size is). This is dense, high-dimensional information, but it is not text yet, and it is not even a probability yet.Here is the part that surprises people the first time they see it: for generating the next token, only the last row matters.The hidden state for the final token in your sequence is the only one that gets used to predict what comes next. Every earlier hidden state did its job already, feeding information forward through attention, but the actual “what word comes next” decision is read off a single vector.The LM HeadThat last hidden state now passes through something called the LM head, which sounds more complicated than it is. The LM head is a single linear layer, nothing more.The math is one matrix multiplication: a [d_model] vector times a [d_model, vocab_size] weight matrix, producing a [vocab_size] output. If your model has a 128,000-token vocabulary, you now have 128,000 numbers, one per possible next token.These numbers are called logits. Two things about them:First, there is no activation function here, no ReLU, no sigmoid; this is a raw linear projection, nothing is squashed or clipped.Second, logits are not probabilities. They do not sum to 1; they can be negative, and a logit of 8 does not mean “80% likely”; it just means “more likely than a token with a logit of 3, by an amount that depends on the whole distribution, not just these two numbers.”Intuitively, a logit is a raw compatibility score between the model’s current understanding of the sequence and one specific candidate next token. Higher means better fit. That is all it means on its own.Logits To ProbabilitiesTo turn logits into something you can actually sample from, they go through softmax:P(token_i) = exp(logit_i) / sum(exp(logit_j) for all j in vocabulary)Softmax does two things at once: it makes every value positive (through the exponential), and it normalizes everything so the whole vocabulary’s probabilities sum to exactly 1.Why softmax and not some simpler normalization, like just dividing each logit by the total? Because the exponential is what gives softmax its sharpening behavior. Small gaps between logits turn into small probability differences, but large gaps get stretched out disproportionately.A logit gap of 2 between the top token and the next one produces a very different probability split than a gap of 8, even though both are “the top token is ahead.” This sharpening effect is exactly why temperature, which we get to shortly, works by manipulating the logits before this step, not after.SamplingWhere exactly does sampling happen?Logits --> Temperature --> Softmax ---> Probabilities --> Sampling Happens After softmax, you have a full probability distribution over your entire vocabulary. Now the model has to actually pick one token, and this is where sampling happens.Greedy decoding just takes the single highest probability token, every time. It sounds like the obvious choice, but it produces oddly repetitive, flat text in practice.Once a model greedily commits to a slightly generic phrase, it tends to keep digging that same groove, because the highest probability continuation of a generic phrase is often another generic phrase.Probabilities --> argmax() ---> just pick highestSampling instead treats the probability distribution as an actual distribution, and draws a token randomly, weighted by those probabilities. This is exactly where randomness enters the pipeline, and it is why the same prompt to the same model can give you different answers on different calls.The hidden states, the logits, all of that is fully deterministic given the same input. The sampling step, if temperature is above 0, is the only place chance gets introduced.Probabilities --> Top K filter --> Top P filter --> random draw from distributionTemperature, Top-K, Top-PThis is the part people actually tune, so let’s get specific.Temperature scales the logits before softmax, not after: logit_i / temperature. A temperature below 1 sharpens the distribution, making the model more confident and more repetitive, pushed toward greedy-like behavior as temperature approaches 0. A temperature above 1 flattens the distribution, giving lower probability tokens a real chance.Top-K is a hard cutoff. Keep only the K highest probability tokens, zero out everything else, then re-normalize and sample from what remains. Simple and predictable, but it has a real blind spot: K is fixed regardless of how confident the model actually is.Top-P, also called nucleus sampling, fixes exactly this blind spot. Instead of a fixed count, you keep the smallest set of top tokens whose cumulative probability crosses a threshold P, say 0.9. When the model is confident, that nucleus might only contain 3 tokens. When the model is uncertain, it might contain 100.For learning these in detail, read below:Generation Control: Mastering AI Output for Better ResultsIn practice, most production systems combine temperature with top-p rather than picking one: temperature shapes the overall confidence of the distribution, top-p then trims the long, low-probability tail before sampling.import torchdef sample_next_token(logits, temperature=0.7, top_p=0.9): scaled_logits = logits / temperature probs = torch.softmax(scaled_logits, dim=-1) sorted_probs, sorted_indices = torch.sort(probs, descending=True) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) # keep tokens up to and including the one that crosses the top_p threshold cutoff_mask = cumulative_probs logits kept as-isInvalid tokens --> logits set to -infinity ⬇️ apply softmax-infinity --> 0 --> prbabality 0So invalid samples can NEVER be sampled!def apply_schema_mask(logits, allowed_token_ids, vocab_size): mask = torch.full((vocab_size,), float("-inf")) mask[allowed_token_ids] = 0.0 return logits + maskThis runs at every token generated, not once at the start. The set of allowed tokens changes constantly as the state machine advances; for example, once you are inside a JSON string value, most tokens are legal, but the moment a closing quote appears, only a comma, closing brace, or closing bracket becomes valid.A few tools implement this in production: Outlines and XGrammar compile a JSON schema or grammar into this kind of state machine automatically, llama.cpp uses its own grammar format called GBNF for the same purpose, lm-format-enforcer does something similar for HuggingFace pipelines, and OpenAI’s response_format with strict mode does this server-side so you never see the mechanism directly.The key insight worth remembering: The model does not know the schema. It has no awareness that it is being constrained. The schema lives entirely outside the model, as logit masking applied at every decoding step. This is also why constrained decoding gives you an actual guarantee, not a best effort. Malformed JSON from an unconstrained model happens on a meaningful percentage of real production requests, constrained decoding removes that failure mode structurally instead of catching it after the fact.Content FilteringSafety and content filtering gets talked about like it is one system. It is actually three separate layers, each catching different things, and none of them alone is sufficient.Layer 1: The input classifier(Pre-model) — Before the model even runs, a separate classifier looks at your prompt and flags anything that matches known harmful patterns. This is cheap, fast, and catches obvious cases before you spend any compute generating a response at all.Layer 2: Baked into the weights (During generation) — This is the layer people underestimate. Through RLHF (reinforcement learning from human feedback) and approaches like Constitutional AI (Anthropic), the model itself is trained to refuse certain requests, not through an external filter, but because refusal became the highest reward behavior during training.Layer 3: The output classifier (Post-generation) — After generation, a separate classifier scores the actual output before it reaches the user, catching cases where the model generated something problematic despite layers 1 and 2, or where a jailbreak got past the earlier stages.The Full PipelineHere is the entire path end to end, in order:prompt → tokens → transformer blocks → hidden states [seq_len, d_model] → LM head (linear layer) → logits [vocab_size] → (temperature scaling, schema masking applied here) → softmax → probabilities → (top-k / top-p filtering applied here) → sampling → output tokenA full end-to-end pipeline flow with each stepConclusionLLM’s output mechanism is still just a probability distribution over a fixed vocabulary, produced by one matrix multiplication and one softmax.Everything else: temperature, top-p, JSON schemas, safety filters, is an engineering decision layered on top of that distribution, not a change to what the model fundamentally does.Once you see the pipeline this clearly, tuning generation stops being guesswork and starts being a series of specific, explainable choices.If you enjoy it and learn something new from this article, give it a clap. It encourages me to write such in-depth technical concepts. If interested in knowing more about topics, read the papers and thank me later.Further ReadingThe Curious Case of Neural Text DegenerationTraining Language Models to Follow Instructions with Human FeedbackConstitutional AI: Harmlessness from AI FeedbackThis story is published on Generative AI. Connect with us on LinkedIn and follow Zeniteq to stay in the loop with the latest AI stories.Subscribe to our newsletter and YouTube channel to stay updated with the latest news and updates on generative AI. Let’s shape the future of AI together!Inside the Final Layer: Logits, Sampling, and Structured Outputs in LLMs was originally published in Generative AI on Medium, where people are continuing the conversation by highlighting and responding to this story.Source: Generative AI Pub — Published — Category: Image AI