Learn to build and master Jev-style decision-making models from scratch using PyTorch and frameworks from Hugging Face. By replacing autoregressive token-by- token decoding with a single direct pre-fill pass, Jev bypasses text generation entirely, providing fast, low-latency typed solutions such as choices, ordinal evaluations, and boolean values. You will learn about inference paradigms and explore RLCD training workflows to obtain honest, well- calibrated probabilities for real-world software logic. Hello and welcome to learning Jev from scratch using only PyTorch and basic libraries like Hugging Face. We will use pre-trained models as a foundation, but try to learn the first principles of inference and learning to understand how Jev works and have an idea of this LLM working paradigm. So, in this video, we'll look at Jev and RLCD, reinforcement learning for calibrated solutions. I hope you like it. The first thing we want to look at is what Jev is and what RLCD is. We will probably move on to RLCD towards the end of this series, as I want to first cover inference and the different paradigms under which it is performed, and then work backwards towards understanding the learning workflows. So, Jev is essentially very simple. If you understand how an LLM works, it has two stages. Pre-fill and decoding. Pre-fill is the stage where the model reads the entire request, it is limited by computing resources, and it is then that the generation of the first token is prepared. Decoding is the generation of a response one token at a time. If you have 500 words that you input as tokens for the task context, and the model tries to generate an answer: first the first token is created , it is fed back along with the 501st token, the model outputs the 502nd; all 502 tokens are fed back, the 503rd is output, and so on. Jev tries to skip the decoding stage entirely and reduce the problem formulation to one direct pass or one pre-fill stage. This obviously has many advantages. Your latency will be fully controlled and your output will be predictable. This means that you specify in advance what you want from the model in the form of choices or boolean values, and it returns exactly what you need, instead of you struggling with syntax trying to make a large language model predictable. So, two things: one direct pass without decoding and predicted results. These are two desirable properties that can be named for Jev. Now, moving on to the Jev interface, you can ask questions in three parts. So, first, we have a choice where you can act almost like you would answer a multiple choice question. That is, you have option A, option B, or option C. Evaluation, on the other hand, does the same thing, but is ordinal in nature. So, you could say that this is a good example where you have nominal options and the qualitative feature is ordinal options. What do you mean by that? So, if you ask about directions: east, west, north, south—they are nominal in nature. East is not greater than west, not greater than south, and so on. So it is impossible to maintain any order within the limits of choice. The options simply exist independently of each other. They cannot be lined up in any particular order. Ordinal data, on the other hand, can be arranged in a specific order. So , what is the customer's attitude? Look, he is calm, irritated or very angry. Therefore, they can be ranked by degree of dissatisfaction, so to speak. So, the APIs provided by Jev for selection and evaluation are different. And finally, there is null, which is simply binary. That is, you ask a "yes" or "no" question and get the corresponding probabilities. I wonder how it's worded. So, nominal data is slightly different from ordinal data, and ordinal data is significantly different from binary data. We will try to create a framework that covers them all, but also consider and respect the nuances that these three APIs have. We will try to form a common basis for everything we do, and on this basis build mechanisms for distinguishing nuances. So, for our case, we will build all three based on one softmax per question, but I will definitely try to show how they can differ. Now let's move on to what prerequisites are necessary for this to be implemented. This won't suit everyone. So, if you have a good understanding of attention concepts , how KV cache works, what logit is, what softmax is, and you are fluent in PyTorch, I think this might be for you. If you're not familiar with any of these things, you'll have a little trouble following this video. Something that would be nice to have, but is not mandatory, is a little knowledge of Hugging Face. We will use the VS Code debugger to step through some of the concepts we will try to learn. It is also very important to know the modern aspects of transformers, encoder, decoder, etc. So, these are really necessary prerequisites . If you don't know them, this video probably won't be of much use to you. And if you know them too—that's great, but not necessary. OK. What is definitely not mandatory is having your own graphics processor . The programming in this course will be mostly focused on inference. We will talk a lot about theory, but in implementations we will stick to the inference stage. So you don't need a GPU for this. You don't need experience in model training or any knowledge of mathematics other than an understanding of the softmax function. That's all. So, let's begin. Here's what we 'll do. We will divide this entire process into three stages. The first stage is regular generation using LLM. We will try to write a query that generates JSON token by token. At this stage, we will look at the prefill and decoding stages together. This will give you a good understanding of the problem and context for further work. After that, we will consider an encoder-only implementation. And then we will look at the implementation based on just the decoder. We will go back to the learning process to understand how it works, and then compare approaches to decide which one to apply and when. Now, before we get into how Jev works, it's worth reviewing the standard decoding paradigm, or how LLMs typically perform inference. This will help us better understand where we are going next and how to change all of this to meet Jev's requirements. So, we have a tokenizer here and a model that we load right away. Again, we use pre-trained LLMs . We will use previously taught LLMs throughout the course. In this case, if you look at the bootstrap function , the model here is a quantized version with 3.6 billion parameters. It is small enough to fit on most laptops. So, we have the model loaded. And now we will check what a regular query looks like. So, we have a tokenizer here. The context looks like this. So, the card ending in 4412 made nine purchases in 3 minutes in four countries. The cardholder called and said the card was still in his pocket. Looks like a fraudulent transaction. And the schema you'll be looking for looks something like this, and it'll usually be your Pydantic object that you'll pass to the LLM provider, for example, through an API like OpenAI. So now the schema contains something called " priority". So let's take a look at this. This is an enumerated type, where the choices are "high" and " low". And then it contains a few more points: a decision, where again the choices are “ fraud”, “ check for fraud”, or “ clear”. And also " escalation". So you see these three have been carefully selected to represent the three interfaces that Jev provides. On the one hand, you have a priority that can be high or low. This is the Jev evaluation API . There is a solution, which again is subjective. And then from there there is a logical meaning. So, these are the three questions I tried to formulate to mimic all the interfaces that Jev provides. We'll then try to break this down to simpler examples to see how Jev works later, but let's first figure out what it does . We go through a primitive query function that simply takes all of this and converts it into a format ready for output. I'll skip a bunch of these data operations because they're not that interesting. So, when this processor comes back , we'll just look at the result and be able to clearly see what exactly was done "under the hood." We will try to clarify this. So, if I look at the request, I tried to put it in a format that LLM providers like. So if I were to type this, there 's an IM start token there. And from there the system message begins, that you are a precise data extraction system, blah- blah-blah. And then there's the JSON schema that I want you to follow. And then the user provides context. This is what Gemma calls a state, and then it's a user request that will be entered into the state, and the questions will come later. Yes? So in this case, in normal LLM decoding, all I'm asking for is: here's the user request, here's the system request, followed by the user request. Can you please give me back the JSON? Right? So, we even open this curly brace so that LLM can continue to decode accurately. OK. We did a lot of work on changing the data format, but it wasn't interesting to us in our quest to understand Jeff, so I skipped a bunch of those processes. So, in this case, we're just tokenizing this entire prompt. So you will get a bunch of values that point to all the token IDs of the corresponding tokens. Of course? So let's move on. If there is a sequence termination token— great, if not—we'll set ours, and it will be an unescaped JSON brace again to indicate that we're trying to finish decoding the generated JSON. Again, I would say that these are all just purely technical points. This is where the actual decoding begins. So let's try to clear up some space to better understand this. Let's take a closer look. Our input data, after tokenization, is passed through a model that produces a result; we take its logits, extract the past key values—this will become our KV cache, and then we look at the last token, go through its dictionary, and choose argmax, which is essentially the token with the highest probability. Now, if it matches one of our stop sequences, we will break. Otherwise, we'll continue and it will all become very clear once I generate this and show you step by step what's happening. So, if we take a model , pass input data through it, we are constantly serving the cache. In this case, the cache is set to "none", but you see that the cache is then refreshed. What we can do here: If we have the same state or context and we want to answer a series of questions, we can cache the context state, and after I ask the first question, I just get that cache and process it by adding the question to that context. If I have a second question, I can retrieve the context from the cache again, attach the second question, and so on. So I just try to maintain the cache, which is best practice. OK. So, from there, uh, logits come out. So, just to understand, I had, uh, let me take a look at the input. So, we have the input IDs, which will give me an idea of the number of tokens. So if I take the length of this, I get one because it's nested JSON. More precisely, it's a nested array, sorry. So, in total, there are about 130 tokens in this case right now . OK. So from there, once the result is obtained, you will see that if I turn to logits, I will have 130 outputs, and each token position will have a number of options to choose from equal to the size of the dictionary. So, uh, the principle of the transformer is this: if you input five tokens, you get five predictions at the output. Each prediction will contain the full probability distribution for the corresponding dictionary. So if you're looking for just the last token, you discard everything except the last prediction, and then use that probability distribution to choose argmax, which is what we're doing here. So, we refresh the cache. Now we take the first measurements from these logits. Since we are only processing one package, if we just take the first one and check its shape, you will notice that the unit in front has disappeared. And what we're left with is 130 by 152 thousand, uh, elements that we want to look at. And then if I take the last element of this, I'm essentially left with the probability distribution of the predictions of the last token. So if I look at the form of this, I'll have 152 thousand outputs to work with. Now I just need to convert them to floating point numbers, run them through softmax, and it will just be dimension minus one . So, uh, it does softmax on all of these numbers. And then I'll take arg max. So if I actually take this and do them one by one on my own, you get a tensor with 130, uh, or rather 152 thousand almost, uh, probabilities. And then if you select arg max from them, you will get the token ID that has the maximum probability. So, uh , arg max gives you a position, which is nothing more than a token identifier. So I can actually decode this and get my answer. Let's see. This is 330. So this is a token that was, uh, which is—so this is a token that was predicted in this cycle. So if I let this run for a few iterations and check the marker that was just created again, you'll see that it's constantly changing over time. Now it becomes a priority. So, I think we'll probably get the closing quotation marks along with the colon. And if we continue to do this for a few iterations, you'll see that the accumulation that happens in this variable will reflect the answer that we're looking for. So that's why you see the ChatGPT web interface giving you tokens one by one. This happens because it iterates through all the previous context and generates these tokens one by one . So, he believes that the priority of this particular ticket or context is high. And if we continue to watch, going through a few of these steps, you'll see that eventually it will reach a point where it will almost generate all of the JSON. So if I look at it again, it has completed predicting the second element. So, he thinks it's a scam. And then we are left with escalation. has more markers to generate. And if we get to the end, after it completes its loop, we can see the final text that it predicted after all the iterations of generating these tokens one by one. So if we look at the text here, you'll see that it returned our final JSON. Let's actually bring it up on the screen so we can see it clearly. We will print the text so that it is displayed in a convenient format. So you can see that starting with this quote for the word "priority" and then selecting each subsequent marker for that text and generating one by one took a lot of time, which could have been avoided if we had formulated the problem statement wisely, which would have allowed us to make predictable assumptions without getting into this vicious circle of decoding, at least for this problem. Yes. So, there are a bunch of helper operations that make sure the result is correct JSON, checking for escaped characters, start, end, etc. And finally, after this post- processing is complete, you get the answer. Now we have seen how the decoder model was used to generate the response. Now we will start with one of the first formulations of our generative model. And this will be implemented through the encoder bus. So, we will try to use an encoder model like BERT or any other variant of it to reformulate our case in a way that avoids decoding entirely and still tries to get the desired results. So, first, we will take a model consisting of only the encoder. We won't do any training and will try to see how well we can adapt it for this specific case. And if not, then let's consider what approaches to learning can be applied. OK. So, first of all , you'll see that we'll code the question first. This is a rather unique approach. Typically, you provide context first and then ask questions. Here we do the opposite . We start with a question, and then we add context here. This approach is chosen because the attention mechanism usually places a lot of weight on the beginning of tokens in the attention block. So we want to make sure we take that into account. Experiments seem to have been conducted in the opposite order , but this option proved to be more effective. Therefore, we will stick with this implementation. Now that the question has been asked, you have all these answer options. In our case, it was " high" or "low ", "fraud" or " verification", or even "yes" or "no". In this case, we will add a mask token before each option. This is a special type of token that we'll insert, and we'll sort of sequence it all together. We will fully coordinate this process. And then we will effectively use this prediction to determine which variant we can predict. So, if you have option B, we will insert the mask token opposite or immediately before option B. If there is option C, we will put the mask token before it, and so on. And finally, all the context will be located here. So as you can see, we'll have to do this for every question. So if you have a very long context, every time you ask a new question, the state has to be resubmitted to this whole system. You may ask: is this the encoder model? Why can't the state be served first? Why can't I cache it? And why can't I just change the questions one by one? Well, that's the problem. Encoders are bidirectional by nature. So, every time you change a question, you completely change the attention weights in each position of your attention matrix. Because of this, even if you change this order and put the state first and the question last, every time you change the question, the weights of the bidirectional states will change, and you will have to recalculate everything again and again. Um, I hope you can see that. It's a very simple interface, but you can clearly see that for each question you will have to keep adding state over and over again . So if you have a lot of tokens and a lot of questions, it might not be very scalable. But, given that it is bidirectional attention, it will give you the maximum return on each parameter of the model. Yes ? So without further ado, let's dive into the implementation of how pure encoding.py works for us. And we'll set a breakpoint here. And we will try it in action. Uh , what we're interested in is that we'll have a state. We'll have two questions to see how this interacts with the multiple-question setup. And then we'll see how we embed the question, the mask, and the context into one. Now, what I've deliberately avoided here is batch processing questions, although you can do that. Uh, I used the old-fashioned approach with for loops so we could break it down and see how it works. This is not the most efficient implementation. Good. So, first, there is a function called build that takes a state, which is nothing more than a context. So we'll take a quick look. So, this is our context. Ticket 8812, double charge for the same order, no response for 4 days, cancel this and refund my money, otherwise I will file a complaint with the bank. Good. So, the condition along with the instructions, so each instruction will be a question. That's how we phrase it. So this for loop is now processing the first question. Uh, once he gets stuck again, there will be a second question. And then there are options here. Again , this is a yes or no question. There may also be several options. So, with that setup, let's go inside the build function to see what's going on. Now it's very simple. We take the instructions, i.e. the questions, and first encode them. So, let's go. We get the token IDs of these questions back. And we do this by putting the classification token first, then the question, and then the separation token. So, uh, once we do that, we'll actually keep track of these separators and delimiters that we're inserting so that we can isolate what we need to find the answer. So here we have a range, a range of questions. And then we have markers. These are nothing more than answer options. So, we already know that for the first question, the options are "no" and then "yes." So, uh, we're putting in a place where I'm going to mark my mask token so I can track it during the withdrawal. You'll see all of this in a minute, uh, if it's unclear. OK. So, we added markers to the identifiers. Now, I think if we track the IDs, we'll be fine. So the identifiers were only related to the question, but notice that we, uh, inserted a classification token and then a separation token after the question ended. A mask token that we will observe and that we will insert before each option so that we can track the probabilities of the masked tokens and perform certain actions on them to assess the weight of that option or the probability or propensity for that option for that question. Yes. So, once we go through the cycle again, I expect there will be another mask token, followed by a " yes" option. So let's see this in action. Right. Once this is added , you will see that another mask has appeared , which is now responsible for the "yes" option. Uh, all options are exhausted. The cycle will no longer be executed. And finally, uh, the identifiers are generated and returned. So if I do a final check, you'll see a classification separator, followed by a question, then a separator, a mask for the first option, a second mask for the second option, then a separator, and finally I add all the context, again with a separator. So, this is very encoder-specific terminology: BERT- style, MLM-style, learning paradigms based on masked language modeling. Uh, and this is used all the time during inference if you're working with BERT. This order of tokens allows us to distinguish what to look at and what not to look at. OK. Uh, so this is the result of our build function. And then we get the result, and it's literally the LLM challenge. So, we pass all our incoming IDs. Now the input ID here contains the question, the option, the state, everything. Next, I extract the last hidden state and transfer it to the CPU. Now we'll look at that in a second. So, my model's prediction is complete. This is what it looks like. So, the form of the hidden state is 54 by 768. This literally means that each token has its own embedding, which we got. If you look at the number of tokens, there were 54. So the hidden state returns 54 embeddings for each of these tokens . Now, if we ask the last hidden state what it will give us... let's run all of this. If I run this, it will do another model prediction. So the result is 1, 54, 768. Now if I take the first element, it will get rid of the package dimension and give me 54 by 768. If I check the form , it is 54 by 768. After that, I convert it to float format. And now, in the next step, I want to extract just the query. Again, the ranges and all these variables were created in advance. We wrote a script for many of them to save time. Let's look at them now. So if you look at the range, it's a tuple. The first element of the range denotes the starting token or starting position of the first question. And the nine here represents the final position of the first question. In fact, if I do 1 to 9 and try to decode it, I get the exact question. Is the customer asking for a refund? Yes. Why did we start with one? If I were to start from scratch, I would come across a classification token that I want to avoid. The question itself does not contain this token. Right? So, I will take the hidden states of the questions themselves. If I look at it, there were eight tokens in the question, so 768, and I'm averaging them. I am trying to create a 768 vector representation for the query itself. Here's how it should be understood. Now, if I take that and find the average, you'll see that I get 768 numbers . I will repeat that. This is the best representation of the question we have from the model. Because we take all the tokens in a query or question and then average their embeddings to get the best representation of the question itself. OK. That's where we got our markers, so we know for sure that, at 10th place, again , we can check that. We have a mask token, and if we look at the token before it, you'll see that, or better yet , if we look at the token after it, you'll see that it's followed by the option that we're interested in. Now, if I look at the 12th marker, you'll see that this is the marker for the " yes" option. So we have two options, one um, one mask for each of the options. And then we take and find hidden states for them as well. Again, we will get these two representations of the embeddings for the two options. Now, this is where our implementation will really be basic, and this is where a lot of the variation in implementation occurs. So, let's try to understand this better. So, we have a question here. And then um mask token, then option one, then another mask token, then option two, and so on until there's a separator that gives me the whole context or state. And then, finally, it gives a forecast. So what we're trying to do now is each of these tokens generates a 768 uh embedding vector. Now, the question, if I take a range of questions, the question itself will span multiple tokens, and I average them. Then this will give me a clean representation of 768 embeddings. Now this mask token will also have a representation , so M1, and this will also have another representation, like this . So you see the range of options is 2 by 768, and the query covers 768 by 1. So that's an option. This is a request. Now, one way to look at this is that I can make them interact with each other. So I can literally take the dot product , and I'm left with two scalars, and I can do something on top of those scalars, even apply a simple soft max to those scalars and consider them probabilities. This is literally the simplest way to phrase Jeff's problem. Um, regardless of whether you have a choice, a zero value, or a score, you take all these tokens and make them interact with each other. It's still bidirectional attention, and then from there you take the query vector, you create a representation of it, you take the masked vectors of options, you make them interact, and you can interpret the output as probabilities. Now, obviously, we are taking a ready-made base model without any training. It won't work, but it serves as a foundation. So, what else can be done about it? Now you have the first option - a linear "head". What I mean by this is instead of taking these 768 vectors for each of the options, I can pass each of them through a linear " head" and get a scalar at the output. So , this linear "head" can literally turn 768 values into one. So if we input a variant vector, we get the corresponding probability at the output. So, this matrix of options can literally be run through a linear "head" and produce two scalars that represent the corresponding probabilities for each option. Second: I can do a full-scale retraining, and I can set a goal, either by taking the way I formulated the problem, where the variant tried to interact with the query, or I can add linear " heads" and do a full retraining. Option three: I do a complete retraining of the base model. In practice, it turned out that the learned embeddings, all 768 of these representations, are actually quite high-quality and highly accurate, but they don't "know" that they need to be output as probabilities . So any numbers you get from them will not make sense, even though all the necessary signals are already present. So it really depends on the complexity of your task, how much data you have or how much you're willing to put into it. You choose one of these options. In its simplest state, you can get probabilities, but that wo n't get you very far. You may need to add a linear "head" and retrain just that. You may need to do a full retraining or full retraining of the base model. So, from that perspective, we come back here. Now we have vectors of options. Remember, it's 2 by 768. And then we have the query vector, which is 768 by 1. So you make them interact and normalize. You actually run them through a softmax and get probabilities. So in this case, if I take a look, you'll see that two scalars are produced, and they can be thought of as probabilities of the answer to this particular question. Now I know they are not very accurate. This should be "yes", although it only gives 30%, but note, the embedding signals are very high quality. This interaction takes some getting used to. A little tweaking is needed , and this is where our three options come in handy. Yes? So , I think we've finished looking at the encoder representation of Jeff's work. I hope this has given you an understanding of what you can do and what you need to do to make it work. I hope this helps. OK. Let's see what's happening in the world of decoders. It turns out that you can also use a decoder to formulate the problem so that only one direct pass, just a fill step, without a decode step, is needed to achieve all of Jeff's goals. And here's how it works. We're going to pack the whole state in advance. This is the exact opposite of what we did with encoders, because this time it can be cached. So if the state, and since this is a causal mask, if there is a KV cache for the state, I can just change the question here and reuse the KV cache as well. Now, one more thing I can do: if I have a large enough context length, I can pack as many questions and answer options here as I want. Imagine you have 32k of context and the state only consumes 1k. Now imagine you have 31k of context to fill in as many questions and options as possible. It's pretty obvious that Jeff is using the decoder approach, as it supports up to 5000 questions for a single state. So I assume they may be using decoding, but again, they haven't released anything so you can never be sure. The first token or approach that we will apply here is that we will formulate the problem like this: here is the context, here is the question, and we will just leave out the colon. This effectively turns the LLM into almost an autocomplete task. And once it issues a token, we'll look at that token, and it will, of course, have a whole probability distribution of its dictionary. We will actually only focus on those tokens that are in our pre-trained or, more precisely, pre-built versions . So, if the choice consists only of " yes" and "no", we will only consider the probabilities of the "yes" and "no" tokens in the first stage of decoding. And we'll ignore everything else, treat them as probabilities, and return them. Now, instead, you can frame the problem differently where you, uh, reframe the problem almost as a multiple-choice question . You have the options A, B, C, and D. And then you ask the language model to answer what the answer is, and nothing else, just A, B, C, or D. Now, depending on the letter you choose, uh, you output the corresponding probability. These two approaches require no training. Everything else is needed. So, uh, we'll stick with that part, and it can be extrapolated pretty easily because, again , it's just prompting the model to output one character that is one of your options. And then you can have, again , a pointer head for the option selection that goes through the linear layer to create an interpreted probability distribution. This is what Jeff implemented. This is an open source alternative to Jeff, and then Cloudflare released Clef, which is essentially a combined range implementation of the entire variant plus a separate line head. So , it takes into account the entire option for these questions, so it will never make a mistake in case of duplicates. So, to clarify what I meant by the last part, is this legitimate, let's say, as a question ? I just made that up. Option one is fraud, and option two is a review of the fraud. Yes? So, if you have a perfectly formulated Jeff-style problem where there is option one and option two, notice that the first tokens are exactly the same. This is where our simple implementation fails. It won't know the difference between them, since the first tokens are literally the same, and it will probably give a 50/50. So, in my implementation, I made it so that if there are duplicate tokens or identical tokens, we simply give an error, prompting the user to change the way they formed the variant. So it could just be a "review" instead of a "scam review". And then it's given for free, right? So, again , we don't have that problem. If you just call them A and B and then ask LLM to provide one of those letters, I think you'll be fine. But pooled span takes this to an extreme and avoids the problem entirely; even if there are exactly the same options, this particular implementation allows you to distinguish between them. We will again not consider any of these three implementations. We will focus only on the first token task and mainly on how the cache can be reused and how question replacement occurs. This in itself will be a challenge. Before we dive into the code, I want you to understand how we're going to form this causal mask block, because it's really important to visualize to understand how we can package multiple questions into one sequence. So, let's say we have a question. Uh, we have some context . Uh-huh. My name is Vivek . Now the question becomes, "Oh, what's my name?" And then maybe, uh. Which, or better yet, we'll put it this way . Uh, an option might be uh Vivek or something else. And "What is my last name?" And it could be anything. So, option A and option B. Okay. So the way we formulate the problem is this. This is, after all, a cause-and- effect mask. Therefore. Uh, token one, token two, and so on. So if you really visualize the attention matrix here, it's like this. The first token can only pay attention to itself and cannot pay attention to anything else. Now the second token can pay attention to itself and the previous token, and so on. This continues, right? So, uh, at each step, each token can only look back, not forward. Now this in the context of having multiple questions and wanting to package them all into a specific context becomes a bit tricky. OK. So, let's say the question consists of three tokens, and let's say , for the sake of argument, or rather, the state is three tokens. For example, you can ask one question along with answer options in the next two tokens. And the second question also takes up two tokens. Again , so we can see the full matrix. So what happens is you have state one, state two, state three, and then Q1—Q11, Q12, and then Q21, Q22. And the same thing happens here: S1, S2 , S3, Q11, Q12, Q21, and Q22. Now, the way you want to package this is that it's definitely going to be regular attention, as you know it, typical masked attention, a causal mask. Yes? Now this question must see the whole context. So I have no problem with continuing the causal mask we are building here. Right? So far, so good. Now, what should happen to the second question? Ideally, the second question should have full access to the state, but why should the second question have access to any of the content of the first question? I mean, it's theoretically possible to give the second question access to the context of the first question, but think about what you're doing. You are effectively poisoning the context with an issue that you shouldn't even know about or look at. This will result in the fact that if you ask the questions in a certain order, you will get different answers. So, it won't be very reproducible in your area. Interestingly, as far as I've learned, Cloudflare's Clef did exactly this kind of implementation where they didn't use masking, and that does n't make any sense to me at all. Because once you break the order of the questions, you will get different results, I am sure. So, if you look at it this way, the first question, or rather the second question, should have full access to the previous state, not have access to the first question, and of course have access to its own token. So, it has access to the state, it doesn't have access to the question, it has access to itself. And here's how we should do it. The state will be a separate block, and when it's the question's turn, it will have access to its own block and state. And the next question will have access to its own block and to the state. And this will be completely disguised. So creating this mask will require some array manipulation, which we'll cover shortly. And that's why I wanted to focus on the basic principles and key aspects of building this, instead of looking at a bunch of different implementations and confusing myself. Once you understand this foundation, I think the different implementation variations, the different approaches that different providers use, will become readily available to you . Without further ado, let's dive right into the implementation. So, we have a clean decoding script here. And since we have some theoretical understanding, we'll take a quick look at the state. Here is the context you will be given. So, "I just want my money back twice ." And there are, in fact, two questions. If I look at the first one: is the customer asking for a refund, yes or no? And which team is responsible for this? Is this the billing department or the technical department? These are two questions that interest us. Again, this can be easily extrapolated to an evaluation question. In this example, I only included the no value and the choice options. Yes. So, first we loaded our model. This is a decoder-only model, so we will use the same model as in the regular decoding stage. So this would be a Qwen model with 3.6 billion parameters. And then we provide the token, model, state so we can get our cache. And this is one of the most desirable features of decoder-only models. I can pre- cache the state and then continue answering questions. I can either combine these questions together or just answer them one by one in a loop. Yes. So, we'll dive a little deeper into the profiling stage to see what's going on. And then you'll see that the whole state is tokenized like this . And then I actually do the model inference here. And output with past key values to get my KV cache for that particular state. So, again, if I look at the IDs that came in, there are 38 tokens. So, a KV cache, as I understand it, will have certain layers and certain KV heads. So, all of this together, and of course, key values, so there will be two of them. So, two times the number of layers plus the number times the number of heads, times the dimension of the embedding. All of this will be important when it comes to extracting the KV cache. So, the dimension, if I look at Oh. This is Let me look at the length. So, this model has 28 layers. This tells us something . OK. And then each layer has this KV cache. More specifically, how do we access the first value? I think we do this because of the nature of layers. This is the first layer. And here I can look and find this specific... Well, here, please. So, here we have keys and values in the form of tensors. So if I look at this form I'll get the batch size. So, there are eight KV goals, 38 tokens, and 128 dimensions. And there are 28 layers in total. So, multiply this by 28 and by two, because these are just keys, we also have values. OK. So, from here we pre-cache our state, so to speak, and are now ready to answer questions. So we're going from Let's run this to this place, and then we're going to take a deeper look at this feature—responses. So, we've figured out the condition or context. Now we are going to deal with the questions. So, we have a question here. This is a list of tuples. Each tuple is a question and a list of answers. Or rather , a list of options. OK. So, here we start the whole process of trying to make sure the right tokens are getting the right tokens. That's all, so that the second question can't see the first question, and so on. So, we create this question mask. Let's take a quick look. So, we have the width of the segments. This will contain the boundaries of the questions themselves. So if you see, if I have a total of 35 tokens that make up the question. So, the state is already cached. We don't need to worry about that . These are the 35 tokens that contain two questions along with their options. And we saw that about half of them were the first question, and the other half were the second question. So, the segments denote the actual index of the question. Yes. Let's convert them to a torch tensor. Now we build a basic causal mask using a-range. And then we make it greater than or equal to its own transpose matrix so that we can get our causal mask. So, if we take a quick look, the causal mask is simply: the first element is " true," the first two elements of the second line are "true," the first three elements of the third line are " true," and so on. This is a common causal mask. Well, we'll need to do more, as we'll see, to get that question- to-question abstraction that we want to achieve. So whether it's the same question or not is determined by the segments. And we have segments in the torch tensor called seg. So we'll do the same thing, but with an equality condition before its transpose. And that way I can assess whether two tokens are part of the same question or not. So if I just look at the tensor of identical questions that I have, let's see what it comes up with. You'll see that the first few tokens here are all true about each other, and the last group of tokens here are also true about each other. So there is a clear division: these are the tokens of the first question, and these are the tokens of the second question. So now, if I take this and make it interact with the causal mask that we created in the previous steps, interesting properties emerge. So we'll just wait for it to happen. Yes? So, in the next step, visibility should be the concatenation of torch.ones, width of state length, along with the interaction of the same question and the causal mask we constructed. This will give us the desired state of what we were striving to achieve. And I think it gets in the way of questions. Okay, so we have a visible tensor. Let's take a deeper look. We had a causal mask, we had a mask of identical questions, we made them interact to get it. Good. So we see that I think this whole shape is 35 by 73. So there are 73 tokens total, which make up the state and all the questions. They were packed together. Now, since the first few tokens are already cached, I only need to worry about the 35 tokens that make up the question. So if I look here, what I have is this whole piece of the puzzle. So, I have all the questions. I have this part. So it's 35 by 73. I don't have the top part because I already cached it. So, this part is already in the cache. I'm only interested in this part. And that's why I have this fragment. I hope that's clear. So, once you see the visible, you will notice that they resemble a state. The first part is the state where everything is true. Again , if you look at this part, it is all completely true. And then the only interesting part of decoding starts when it crosses the number of states. So if I really wanted to split this up, to ignore maybe the first 38 tensors that contain context. So, let's really do it. So, starting with the 38th token, you'll already see that we have this causal mask working, because we've exhausted the state that had 38 tokens in it. And then from there you'll see that everything on the bottom left is false because those are tokens for the first question, and we're working on the second question now. And then everything is simply causal for the second question. Here everything is causal for the first question. Everything else is false in the upper right corner. Everything else is false in the lower left corner. So, we are trying to maintain that mask that we saw here. That neat mask where everything has to be turned on. Everything here should be turned off . And everything here should also be turned off and so on , right? So, we only look at questions that consider their condition. This is what we want to achieve. I hope it was clear how we build this mask. It is slightly different from the usual causal mask of attention. Therefore, we had to delve deeper into this issue. But once we get the mask back, everything else remains the same. We will follow the same path as in the normal decoding state, where we make only one prediction and nothing more. So, we have all these segmentations and more . We will pass that attention mask that we built, along with the identifiers containing all the state, questions, and so on. We'll pass our KV cache here because we already did that for the question. We will also pass the position IDs here. And if you look here, you'll see that they cover the range from 38 to 53. So, again, keep in mind that while we're serving the new tensor, the cache of 38 tokens has already been served, so the correct positions had to be filled in so that the positional embeddings, rope and everything else could be calculated correctly. Yes? So, we get logits. And we'll just let it run for now. So we get logits for this. And if you look at the form of logits, typically LLM should return the same number of tokens that you passed in, but for each token position there will be a dictionary size that will contain the probability distribution for each position. Now we will only be interested in the last token, since we assume that the last token itself contains our answer without duplicates. So we're going to go deep here. So, for each instruction variant, solution IDs in questions, solutions, and candidates, let's take a quick look at what these arrays do. Questions are simply the text of the questions themselves. You've seen this before. Decisions, 18 and 34 are actually the final decision point for each question. So, uh, if I look at the first one, 18, you'll see that it contains a question along with answer choices. And if you look at everything else , you'll see that there's a second question there, as well as its variants. So, decisions are made based on these questions and their options. And, of course, we have candidates. These are token identifiers for "yes" or "no", "billing", and " technical". So we'll also take a quick look at the candidates to be able to assign a probability. So, this is "yes", this is "no", this is " yes". So, 902 is "no" , 9834 is "yes", 30531 is "billing", and 10916 is " technical". This is what it, uh, looks like in the end. So we'll go through this list of questions, solutions, and candidates via zip. Now, obviously, we are only interested in decision points. So we'll move right there. This is where the first question ends, and it was able to take into account the entire previous state. And here, in the next cycle, the second question will end, and it will also take into account the entire state and context up to this point. So this is the last token I'm interested in, whose logits I'll take and convert into probabilities. And in my case, they will represent the corresponding probabilities of the options themselves. So let's take a look at them. So, we have probabilities here. A quick look. So, for the first question , we have 95% confidence that the customer wants their money back. So, as you can see, everything works right away with the decoder . Unlike embeddings, which require some additional training. OK. So for the second question, we'll go the same way: which team is responsible for this ? Is this billing or technical department? And then we will see the probabilities for these two options. And if there are three options, then we get three probabilities. This is not a problem. So, 97% confident that this is a billing team issue. And about 2.5% for the rest, which means the level of confidence in the technical team, right? So we move from here to the end. And practically when we get to this point, you'll see that I have another line of code that says, " Response: How angry is the customer? Calm or angry?" This is a graded question that you can run yourself . I didn't start it. But the magic of the "decoder-only " style that I wanted to show you is that it doesn't require recalculating context or state. They are already pre- calculated. So you can keep adding more questions here. And changing them, changing the options—it won't change anything, it won't affect the result. Everything will remain the same . So, the situation will never change. It is already cached. You can even ask the user additional questions instead of grouping them into packages . I hope this gave you a good understanding of how output works in decoder-only models . As a bonus, we will try to look at how learning works. We won't be looking at code or any datasets and such. However, we will formulate a learning objective. We'll look at the loss function, how it's calculated, what its varieties are, and you'll really get an idea of RLHF in theory. So before we dive into the learning part, I wanted to talk about some of the drawbacks of the method I showed and how they can be easily avoided. When two responses start the same way —and I've shown this before—if two responses start with the same token, our implementation will break. A free way to adapt this to the simple approach we've chosen is to simply rename them. If not, we need to evaluate each token, and this will take us one extra pass. So instead of a single decoding step, we might have to do multiple steps and then apply an averaging strategy . Or you can significantly shorten each option and code them with letters, like A, B, C, D, similar to how tests are worded and answered . This will also eliminate the collision problem. Of course, you've seen this before in encoders , where a masking token is placed directly before the answer option. Since it is positional in nature, it will never confuse the encoder. Duplicate tokens in variants are only a problem for the decoder, not the encoder. Yes. So with that, we'll move on to the training part. We will look at examples of how training works, but we will not show the training process itself with obtaining a finished model. Let's leave this as an exercise. Some additional or bonus scripts that I will share on GitHub will contain these training steps, but they will be with synthetic data. So, they will be of little use for real-world tasks. But I'm sure you can take them as a template and adapt them for your own use. So now we will look at how one step of learning is performed. These learning processes are usually built on the basis of teacher-student distillation. That's all. So, you either have pre- labeled data, or you have a teacher model that is able to produce certain probabilities. Let's say for our previous question: is it a billing problem, a technical problem, an account problem, or something else—you have three options, and the teacher model says 0.7, 0.2, and 0.1. This is a probability distribution. Now our results can be viewed as logits or scalars coming out of a linear layer or something similar. So, they haven't gone through the Softmax function yet. These are our estimates. Initial estimates: +1, +0.8 and +0.20. And our probabilities are 0.44, 0.36, and 0.19 . Note that the ranking order is preserved, but I still would like our answer to be closer to the teacher's answer. So what should we do now? We now use a paradigm called GRPO—group relative policy optimization. This is a concept from reinforcement learning. If you haven't worked with GRPO before, you should better understand its essence. I think we're going to be really distracted if we try to explain what GRPO does, but essentially it tries to track multiple trajectories, find their average, and if the trajectory follows a good path or generates a very accurate response, we subtract the average from them , and what's left is called the advantage. And we use this advantage to build a loss function, and then apply backpropagation of the error so that the LLM produces good answers more often than bad ones. That's what reinforcement learning is all about. In this case, we don't have multiple trajectories because, again, we have limited our model to generating only certain token probabilities. So we'll have to invent them out of thin air . So, we take these estimates. Here are our ratings. We have 1.8.2 and add random deviations. We add random perturbations to these estimates. Suppose the first trajectory with a certain random perturbation gives the following estimates. In the second step, we add another set of random perturbations, and we get this, this, and this. So instead of the original estimates, we modified them by adding random deviations to get these specific results . Now, if we convert them into probabilities, they would look something like this. I like the first trajectory because it shifted towards the teacher while maintaining the rank order. So I will rate her higher than the others, and since each of them gets a score, their average will be subtracted, and you will see that she has a positive advantage. This is exactly what I was trying to explain - advantage . So, anything that has a positive advantage will have a smaller loss function and will be more desirable compared to others. This will push our student model in the right direction , towards mimicking the teacher's distribution, so to speak. Yes? Just to show you some of the calculations. So, here are the rewards we received. -0.33, -0.52, -0.5241, -0.4154, and their average value is -0.4498. Now, if I take the reward and subtract the mean from it , I'm essentially doing -0.3380 minus -0.4498 divided by 0. 0.0901, which is the standard deviation of this entire distribution, and we end up with an answer that I literally call the advantage. So, these two trajectories have a positive advantage, while these two have a negative advantage. And if you look at the answers themselves, you will intuitively understand why they have a positive advantage. So we had 0.44, which is our baseline. He raised it a little bit , while lowering this one, and this one lowered it even more , or maybe this one went up a little bit, so there's an overall advantage in pulling towards that trajectory. Now this option preserves ranking order nicely and also leans more towards the teacher model, so it has more advantages. So, the scoring mechanism was designed in such a way that with these trajectories, better answers have more advantages, which then has a stronger impact on losses. Yes? So the final question that remains is, where did we magically get these estimates from? So let's look at one of them, how we get or calculate one of these estimates. So this estimate, -0.3380, for the first trajectory alone, is a combination of several estimates. So, we have a logarithmic estimate. We have a spherical estimate. It's a combination of them. So, logarithmic estimation tries to answer the question: "Did you place the probability where the teacher placed it?" Spherical—does the entire answer point in the same direction as the teacher's? So, in general, there needs to be a trajectory towards the correct answer , but there also needs to be an approach where if the teacher model rates a certain option highly, the student model should rate it equally highly. So, it is a simulation of distribution on the one hand and a directed sequence on the other. This is the intuitive part. So the way to implement it is that we have the teacher probabilities: 0.7, 0.2, 0.1. We have the student's probabilities: 0.52, 0.31, 0.16. I take the logarithm of this, which gives this result, and then I weight these logarithms by the teacher's probabilities. So, the more likely it is for a particular option , the greater the impact it will have on the total. Now that I have this sum, I'm going to put it as a separate term , and then add a spherical estimate to it. So, how does spherical estimation work? These are simply probabilities again weighted by the teacher's probabilities. So, we're not just taking the logarithm of the probabilities weighted by the teacher's probabilities, we're also taking the literal probabilities weighted by the teacher's probabilities. So, we take 70% of the first probability, the student's probability , and so on, and so on . We normalize all of this. And then the final score is the one we enter here. So, these two things together form the reward for a particular trajectory. Now, here's where the nuances come in. Um, 99% of this is very clear. So, there is a teacher model that gives out probabilities. There is a model student who tries to follow this and emulate it. These are the logarithmic score and the spherical score. But there is one last step , just for one specific interface. Remember, Jev has interfaces with scores, zeros, and choices. Just for scorers, we have another implementation called RPS or ranked probability score. So, we are trying to answer the question of whether a score has equals or is it ordinal in nature. Let's say I have a question where the answers could be, say, " irritated" or " very angry." This is a spectrum that goes from, you can visualize it as moving from left to right. So, let's say the truth here is that the customer is very angry. So it's zero for calm, zero for irritated, and one for very angry. Now, if I say " irritated" versus if I say " calm," you'll notice that saying " irritated" is a little better. I mean, I would say the model is a little better calibrated than the other model that says "calm." Isn't that right? So, this only applies to ordinal predictions. So for this, if I actually predict this as 0.1, 0.8, 0.1, then the log score, or compared to 0.8, 0.1, 0.1. So if I just switch that around, the logarithmic score will give me the exact same result. But you will see that the RPS will be different. Probability score, ranked probability score will be different. So, it's the squared difference between the accumulated amounts, averaged over minus one of them. So at each step, I'm going to take this RPS for the multiple choice questions only and feed it back into my loss function in some way. Yes? So the loss function finally equals the log score, the 75% spherical score, and you can, in principle, weight it however you want, minus the RPS. So this is an estimate, and in order to build the final loss function, I actually saw things like cross-entropy plus this RL loss. So we're going to use the full cross-entropy plus the RL loss, and by RL loss I mean spherical plus ee, the logarithmic exponent. And if it is an evaluative style, then it will also have a built-in RPS. So, that's the basic concept of how RL CD works. So, there is a full cross-entropy, as well as a full RL component. For this, trajectories are generated using GRPO by adding random shifts, and we try to combine the two loss functions to get the final RL loss, and then finally, uh, we do backpropagation of the error on everything we train. You can choose to train through a simple linear head, or through LoRA, or you can do a full-scale backbone retraining. The choice is yours, it depends on your appetites, the accuracy metric you need for your task, and so on. It really depends on the situation. A brief note on the implementations and some of their comparisons, advantages, disadvantages, etc. So, we implemented both paradigms for Jev: encoder-based and decoder-based. Some of the criteria you should use to evaluate whether to choose an encoder or decoder approach are listed below. If the state needs to be stored between questions, the encoder is forced to reread it each time , while the decoder can read it once and cache it . So, if you have a lot of questions about a limited context , I think a decoder is the right approach. Here, the value of questions accumulates, it increases with each question. Here it remains unchanged, only the questions are new. If you need a longer context length, if you have longer documents to process, then a decoder is the right choice. The status is read taking into account the question. Again, since this is a bidirectional implementation that encoders have, it will always have a slight advantage over a causal mask or causal attention. Variants are multiple tokens. Since the token is in front of , it is very easy to insert ee mask tokens in front of these options and ensure they scale well. Again , with a decoder this will be somewhat of a disadvantage compared to encoders. For multiple labels at once, we can easily replace the head with a sigmoid. We're limited to just softmax, so we'll have to do a little bit of uh-uh work on the model engineering, tweak the architecture a little bit to make this happen. Therefore, usually fewer parameters in encoders give better results compared to analogues. So if I have an encoder model with 400 million parameters, I will need a more powerful decoder model to perform the same task. So, this is the main aspect. So, if you have limited memory or are looking for maximum speed, you should consider all of these parameters along with your field and task to make it work. And finally, you need to write one sequence for queries regarding mask and cache positions. These are engineering tradeoffs that you will have to take care of. Yes? I think that sums up everything we've tried to gather here. I hope you enjoyed it and enjoyed learning. Thank you very much for your time.
Generated algorithmically for Search Engine
Indexing.