Taking Control of LLM Outputs: An Introductory Journey into Logits

Speaker 1 [00:07]

Thank you. Hi everyone. Emek here and I'm a fresh graduate of master's program of informatics at Munich and I'm working as a part-time machine learning engineer and I'm also building my free time building Kipu to help Abitur students in Germany. So I'm building an EdTech field so if you are yeah interested in it just hit me up after the talk I have to talk about it but today we will talk about how to control LLM outputs and this is my third PyCon but I'm presenting for the first time so also with your feedback you can find me after the talk and the goal of this presentation is that like you understand what's going on during the text generation process of language models and why the structured generation is a game changer while working with large language models and And also, what are the libraries you can use in your projects. So, let's start. So, the structure is like going from fundamentals to a bit practical implementation. We will shortly talk about theoretical concepts and what are we trying to achieve. Then, we will talk about how to use these concepts in practice. And, yeah, then we will have a look at the structured generation approaches and the libraries and tools that we can use in the projects. And, yeah, hopefully we will have some time for Q&A as well. So we can start by talking shortly on the setup I used and this transformer architecture. All the examples and values I show in this presentation will be outputs of LAMA 3, 3 billion instruct model. But these are independent from model type, so you can also test out with other models. And this can be reproduced. And for this presentation, we accept a transformer model as a black box. and we live answering what's going on in the transformer to mechanistic interpretability guys I think there's a simultaneous session going on explaining the attention mechanism and KV kitchen mechanism but we are more interested in the input and output of the model today and we will be looking at very first and very last step of the text generation process so see I'm sure you had someone asking you how these models generate text that's good and maybe your friend maybe your grandma or someone like this and we all have a good answer for that I guess yeah token by token that's the one works and it has become a cliche and token by token it became too popular that everyone seems to agree on it but once you think about it like it's not telling much so to dig it a little bit more. I will just peel back two critical layers from this black box, input and output embeddings, and also logits appear at the end, but we're not considering them as part of the model, the logits. So we said basics, right? So token, quick definition, building blocks of the text in LLMs, processed in the LLMs, so let's say they are words of the models, and is they are words of the models the collection of them of course the vocabulary and each model family has its own tokenization approach and vocabulary in different sizes and different shapes and for LAMA we have 128,000 tokens in the vocabulary and accepted special tokens of course and then we have the input embedding matrices. So each token is represented with three thousand dimensional vector and these vectors are learned during pre-training training of the model and here in this example as we have five tokens and a vector to represent each of them this is like five times three thousand embedding matrix and this is the input of this deck box model. So this matrix corresponds the model's internal representations of the input of our input then it goes through many attention and MLP layers which is the black box part and we are not interested in that part and in the end we have output embeddings despite having having the same shape here as you can see output embeddings have a different purpose projects if they project model states onto possible next tokens so So taking our transformed representation of our inputs, we perform a dot product with the output embedding metrics and it produces raw logit values. And the logits are actual outputs of the model. So these are not tokens, but the logits are the actual output. And they are deterministic, so every time we get the same logit for the same input and it says the shape for yeah for now it's like 5 to 128k for each token so yeah it looks like this each column represents one of our five input tokens and each row is a potential x token so higher values indicate tokens that model thinks like most likely to follow and these exact numbers will always be produced again and this is where we can interview to control generation process and managing and manipulating these large metrics efficiently is a practical challenge. So during the training these scores are used to calculate loss as well for the backpropagation purposes and they create the signal to update the model weights but for now we we will focus on the last column and it's the one that is used for selecting next token and as you can see the expected continuation of the model has already high score compared to others. So these are all confidence scores as we said they can range from minus infinity to plus infinity and individual values have little meaning in isolation and others. The relative differences between when logits reveal token preferences of the model for the continuation and we can also consider them as like strength of evidence to proceed in that direction. The deterministic calculation, same input, same output and we can start by defining simplest sampler here. The sampling process is the selecting which token to proceed with. So greedy sampler is the simplest form it's just taking the highest logit value and continuing with it so it generates the same output all the time no and no randomness included no fun and it's very efficient because it's just selecting the maximum and it generates some sequence for the same input so no random subtle so this is the simplest form but yeah it's no fun of course uh so we can what we can do is like we can apply softmax normalization to the logit values here and in the end they turn into a probability distribution and the process until now is like still the same still deterministic we get the same probability distribution so everything is good if we don't select the highest probability now we use multinomial sampler and now we have some randomness because like 70% of the time the most likely token will be selected here but 30% of the time the model has alternative continuation paths. But still the user has no control over the randomness of the input so you see where it's going, right? It's like for giving user this ability we add a scalar in this softmax function to the logit values and this scalar we call it temperature actually so you see every logit value is divided by temperature value so that we have practical control over the randomness of the model and in the end with 0.5 the values closer the values lower than 1 actually amplifies the differences in the logit values so it's I mean the raw softmax output so that we have 98% of the like most movable continuation pattern now with 0.5 but if it's like more than one it's it creates like more uniform like distribution and if as it goes to infinity you can see that it's a uniform distribution over token set and as it goes close to zero it becomes similar to VD sampler so we also have like two more approaches actually more but we will talk about two more the top K limits the selection to K most probable tokens it's like I'm sure most of you also used in practice this and it only considers the K highest logic values it's simple but rigid so it's a defined number and we also have top p it's more adaptive and it's dynamic SLX just enough tokens to reach cumulative probability p that is defined and it's more flexible and it's also known as a nucleus sampling and yeah whether to apply softmax first or top k and top p first it appears to be like a design choice because I like checked a few implementations and there were no consensus on which one to apply first but you can define these. These are logit processing steps and then you can apply multinomial sampling to get outputs. There's one more sampling method that I would like to discover today. This is beam search and instead of modifying the probability distribution beam search we can maintain multiple candidate sequences in parallel to proceed and considers a few next token for each candidate sequence, we can define it, this example is for B, beam width of 2 and the advantage of beam search is that it attempts to find the optimal sequence rather than next optimal token You can also see it in some models now, it's like as they're proceeding with some text, they just go back and change the text and continue producing. This is mostly because of the Beam Search algorithm and it's computationally more expensive and it adds a bit of overhead, but they actually optimizes the sequence quality. And now we look at the sequence that has the highest cumulative probability here, we select two, then we generate two alternative continuations for each of them, then we select two best sequences here and eliminate the other ones, then we generate two more, we select the best ones and we proceed like that. So as I said, you can use it for translation or symbolization, but you can check it if If it works for your use case, because it can, yeah, from my experience, it can generate like, for some use cases, it can generate like very good results compared to multinomial sampler. And yeah, it's computationally more expensive and it tends to produce less diverse outputs. And yeah, I want to highlight one more use case that logits are popular for, and it's uncertainty quantification, uncertainty metric or detection. and we accepted that the models make mistakes and they hallucinate but it's still problematic that like there is not a single way of solving those problems but at least we can predict when the model is likely to make a mistake and start producing gibberish information. So another use of logits here is like defining uncertainty metrics and the first technique is simply taking the highest probability after sampling and use it as a confidence score and it works for simple cases but we also have margin of confidence it takes the highest probability and highest two probabilities actually and check the difference between them to assess model certainty and entropy considers the full probability distributions not only the highest probability order top two probabilities but these are all probability based approaches so they apply after softmax and we can apply logit-based uncertain quantification as well and we will not go into much detail here but there are two types of uncertainty, aleatoric and epistemic uncertainty in the research and when it's aleatoric it's like model knows multiple correct answers. So there's evidence in multiple directions to proceed with and that systemic uncertainty is like model is not certain about how to proceed, what to answer. So the strength of evidence is accumulated during training and it is actually reflected into this logit values. So if we check the example here, these logit values and we get the same probability distributions. However, the model's confidence on its answer seemed to be higher in the second case. For such cases, look-at-base approaches are more popular, and they are thought to be more effective. I added two papers on the resources. You can check them in the pre-talks, and you can just find more information there if you're interested. So let's get back to sampling and talk a bit about structure generation. So it's basically like making model follow some criteria, like generating valid JSON or following some rules that we define. It's a game changer. If you think about like if you use structure generation or not, you probably like 90-95% of the time you need, because if you want to use large language models in like production level systems it's really helpful and it eliminates the need for post processing and the outputs are directly compatible with the downstream systems so I strongly suggest you if you haven't checked that. If you use also yeah if you use it in production level system as I said like it's also important and one thing you can do is like you can elaborate on a prompt and and you can cross fingers with send and sometimes you win, sometimes you get structured outputs or like you just need to add one more instruction to the prompt that will be the also the case then in the end it can be something like that. I like this meme, I saw it in another presentation by the writers of outlines library that we will also talk about. You can sample what you want or what you need or you can just try prompt and pray. So we have function calls and JSON models for a while already so these are being used and JSON model ensures that you get valid JSON and how it's performed is again logit manipulation in the end but no model providers also adapts their training process to generate valid JSON so it's not just in the end but it's also somehow have incorporated in the training process. But in general we owe these well-structured outputs to the methods called constraint decoding. So let's look quickly at how it performed. Constraint decoding, it is restricting the output of our language model to follow specific rules or patterns as here. So the example we show is like token masking. There are like more approaches and with rules we manipulate the logits directly or the sampling process itself and these constraints can be defined in multiple ways like lexical constraints require certain words or phrases and structural constraints we enforce like grammars, regular expressions or JSON mode and for grammars you can use them for specific needs but structural constraints, they can be like regular expressions and JSON mods, so these are the examples. And token masking is basically eliminating tokens that do not follow the specified criteria. So when you are generating valid JSON, there are just limited number of tokens that we can proceed with and some of them are not compatible to proceed, so these are just eliminated in the logit scores and then you sample from what is valid for the next step. So when we check, when we want to check it and practice it's like I added OpenAI here but like all model providers have JSONMOD support now and they don't provide logit scores by the way but you can use JSONMOD to generate structured output and they also integrate like automatic parsing so that they support pydantic models. As an alternative you can also do it with function coding and you can create outputs that follows the signature of a function. But I would like to tell you about two libraries that you can use for structure generation so that you don't have to depend on model providers. I will start with my favorite. It is Outlines. It's started by .txt engineers and it's now community supported as well and it was all before those like JSON mode and function callings so it started very early and now it's being used in many production systems and it works smoothly with Transformers and VLLM so you can use to get structured output from local LLMs as well and it works as a logit processor again it manipulates logits and it just eliminates the ones that are not compatible for the continuation and yeah here you can see a simple example for multi-class classification task this way you can make sure that the output is one of these categories categories this is like customer request classification but you can also use for your own use cases and yeah as I said they work with local models so the automatic parsing is also included here so you can use it with Pydantic models to parse the text directly to your data structures and as it's now performing multiple inference processes simultaneously it doesn't add complexity. So it's actually the opposite, it can speed up the generation process because Because if you see the example here in such a JSON schema, it's the labels and braces and columns, these are not going to be generated by the model, these are just appended in the generation process, and only the parts that model needs to generate, and in this example like 101, then wireless Wi-Fi headphones, these are only the ones that will be generated, And it reduces the number of tokens to be generated significantly, and it can speed up the text generation process as well. And with outlines, you can also generate text valid to some regular expressions you define. It's useful for text normalization use cases or giving users some flexibility in interaction, or you can use it for your own use case. And context-free grammars, it's like I'm not an expert on this topic, but I just tested out for this presentation actually and it's cool and you can define grammars for your specific needs it's basically like coding your sampler or logit processor and we are tested out if you have like a specific use case I definitely suggest and the context-free grammars are defined with the structure called Lark, Lark language so you can check it out as well and yeah if you're interested experiment it out and it's also experimental feature in outlines and lastly the samplers we have discussed today are also supported at outlines library so if you want to experiment them directly you can just use outlines and you can test your use case with multiple samplers and see like which one works better for your use case and yeah these as the pro API the model providers do not provide logit values for the inputs so that you can only use it with local models so the sampling process is done on the model provider side for the API uses and guidance is also provides a similar functionality with different users. With SELECT you can run a multi-class classification again and you can use context management and you can query the chat endpoints in your local models again. And there's one other thing that I found cool about Guidance is their token healing approach. So it solves the prompt boundary problem and it occurs when language model tokenization splits the text at the end of a prompt in a way that it biases the completion. So in this text here the input ends with a like colon and the tokenizer doesn't have the token that has double slashes so it just continues with like space double slash and yeah it's not the optimal case so token healing solving it by backing up one token and constraining the first generating token to have like a prefix matching the like last prompt token and it allows more natural and more control structures and guidance supports token healing only for the local models again for as the logit values needed and it's not supported for api but It's something I find cool about this library and there are more. Structured generation has an active community, so you can check and libraries are evolving, improving rapidly. So I suggest you to try some of them and see which one is more suitable for your use case. So, yeah, thanks a lot. I have uploaded all the supplementary material. And, yeah, I'm looking forward to your questions, and if you want to discuss more on edtech or, in general, in structured generation, the rocket values, just feel free to reach me. Thank you. Yeah, Emek, thank you very much. And, yeah, so now we have five minutes of Q&A. Please don't forget that you need to submit on Slido. I'm not going to be walking around and handing out the microphone for time efficiency purposes and because I'm lazy. So, yeah, I will start with a question of mine. Can you describe a use case in your, I don't know, real life or work life where you've been using what you were talking about today? Yes. so we are I will just explain one use case that we are using for evaluating evaluation so we are using raw logit values for loss calculation and we call it actually in the in the work we call it surprise calculation because when you when you think of it these values are how the model thinks how to have to proceed and there is not only one there's not only one way to proceed but it gives a score to each possible continuation so we just take how it's actually the process used in supervised fine-tuning process we just use like how we want what are the like actual ideal continuation paths and we just calculate the loss or the model's confidence and our continuation paths for multiple responses and then use it to not to like back propagation to the model but we use it for as an evaluation metric to see how models is like progressing for our use case for our specific use case and this is the the loss is the surprise actually like how surprising our output to model is and as this surprise factor getting lower and lower again closer to zero it means for us that the model gets amplifying our use case in its weights and it's getting better and we are using for evaluating the prompts we are testing out. We have a question from the audience so she's moderating okay I will not comment on that I'll make an exception I'll walk over and hand over the microphone sorry maybe my email provider my question is a lot of people use LLMs now for classification tasks and do you think there's a way to merge you know similar logits like for example if you do sentiment classification and you just look for the token positive and negative you may want to also consider other positive and negative sentiment logits in the whole space so I was just wondering if you had any thoughts on on this pattern yes so it's you mean in the generation process many tokens that can gives the hint that it's a positive and the money tokens that that can give him to that's a negative sentiment so incorporating it so it's it's a specific use case so it's not like selecting one of these like four use cases but if you think like there are possible continuation lines in the text that gives the hint of like a positive and you can also do not have to like rely on one or two logit values but you can also take like a sequence and incorporate it but it will require for such a specific case it will require like a specific logit processor definition in the end but in general when the like when the models and the models like the logic scores confidence scores is like masked by such an approach is like this basic from like Occam's razor point of view very basic classification can also work and like includes like different sentiments in just yes but it means like strength of evidence from a lot of different information points so this can be also this can also work from like the simplest approach but you can test out with different custom look at processors in this case awesome there are two more questions however unfortunately the 30 minutes are over so um you will find emmick in the break and uh at the pub quiz tonight as well so i'll uh i'll advertise it one more time and thank you emmick very much thank you and thank you guys for listening so yeah a round of applause please