Getting a Foothold in Reinforcement Learning for LLMs

A developer proposes an opinionated approach to learning reinforcement learning (RL) for large language models (LLMs) by framing a familiar supervised learning problem, such as MNIST classification, as an RL problem and swapping cross-entropy loss with a vanilla policy gradient method (REINFORCE). The method, inspired by DeepMind researcher Ian Osband, aims to build intuition for RL concepts without overwhelming newcomers, using code that simulates an environment providing rewards instead of labels.

Premise Fine-tuning LLMs is a key part of applied AI/ML engineering work these days, by using supervised fine-tuning SFT and increasingly in combination with Reinforcement Learning RL methods like GRPO and its variants, aka policy gradient methods. But RL based fine-tuning approaches are harder to get right, there are more hyper-parameters to manage than a typical supervised learning SL setup, it also introduces a lot of new terms and concepts, not all of which are equally important in context of LLMs. Overall it can be hard to figure out where to start without getting lost in various details and corresponding theory. So if you are someone like me, who have been working in ML since pre-LLM era, trained supervised models using sklearn, pytorch etc., maybe fine-tuned embedding models, maybe tried libraries like TRL from Huggingface, and want to learn more about RL in context of LLMs without being overwhelmed, then this might be a relevant article for you. Here I propose an opinionated approach to get started with RL. Approach The main thing that worked for me was to frame a familiar SL problem as a RL problem, thats it. It will likely give you a sub-optimal result but will help with getting a better intuition, and pave the way for further reading, experimentation etc. This was in part inspired by a twitter/x thread https://x.com/IanOsband/status/2034995355037626712?s=20 from DeepMind researcher Ian Osband. So how to do it? Take your favorite toy classification problem, for e.g. MNIST, use labelled data to simulate an environment which can provide a reward instead of labels. And then swap cross-entropy loss with a vanilla policy gradient method loss, aka REINFORCE. I have some code here which implements just that, lets go through it. First, let setup our dataset and basic feed forward network as our model/policy. python import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms from torch.distributions import Categorical train data = datasets.MNIST './data', train=True, download=True, transform=transforms.ToTensor test data = datasets.MNIST './data', train=False, download=True, transform=transforms.ToTensor train loader = DataLoader train data, shuffle=True, batch size=16 test loader = DataLoader test data, batch size=100 class Model nn.Module : def init self, input dim: int, hidden dim: int, output dim: int - None: super . init self.model = nn.Sequential nn.Flatten , nn.Linear input dim, hidden dim , nn.ReLU , nn.Linear hidden dim, output dim def forward self, input : return self.model input Now lets create a function which implements the training loop given some loss function, and other hyper-parameters as input. Also a utility function to calculate accuracy. The details of network and hyper-parameters are not that important for our goal here, there is no specific reason for choosing this configuration, and the code is kept simple for learning purpose. python @torch.no grad def get accuracy model: Model, data loader: DataLoader - float: model.eval correct = 0 for x, y in data loader: correct += model x .argmax dim=1 == y .sum .item return round correct / len data loader.dataset , 4 def train loss fn, epochs: int=10, learning rate: float = 2e-4 : torch.manual seed 42 model = Model 784, 100, 10 optimizer = torch.optim.AdamW params=model.parameters , lr=learning rate for epoch in range epochs : model.train for x,y in train loader: logits = model x loss = loss fn logits, y optimizer.zero grad loss.backward optimizer.step train accuracy = get accuracy model, train loader test accuracy = get accuracy model, test loader print f"After epoch {epoch}, train accuracy: {train accuracy}, test accuracy: {test accuracy}" Ok, so far this was a familiar workflow of training some classifier with labelled data. For the SL we will use pytorch's cross entropy loss function, for RL lets create a custom loss function which implements the vanilla policy gradient method. python def vanilla pg loss input: torch.Tensor, target: torch.Tensor - torch.Tensor: dist = Categorical logits=input index vector = dist.sample logprob vector = dist.log prob index vector reward = torch.where index vector == target, 1.0, -1.0 return - reward logprob vector .mean Now lets think about this loss function and what exactly it is doing. It follows the same function signature as pytorch's functional cross entropy implementation which is taking input tensor and target tensor and returning a scalar loss value as a tensor. But inside the code we are sampling from input instead to taking the greedy approach and picking the index with highest probability. We first convert input logits into a categorical distribution, and then sample an index from it, if it matches to index expected class then we assign a positive reward 1.0 else negative reward -1.0. The act of sampling here introduces randomness and enables exploration. So if cross entropy loss is $-\log p c$ where $p c$ is probability of correct class/label, the policy gradient loss is $-r \cdot \log p s$ where $p s$ is probability of the sampled class, and $r$ is the reward, which is calculated from whether the sampled class is equal to the correct class or not. This minor difference is the key change, we are basically using ground truth labels indirectly to provide learning signal. The reward itself is a constant and have no gradient here, same for the sampling part, its non-differentiable. During backward pass the gradients will only flow for the $\log p s$ part. Now, lets run the training using both loss functions one by one, it should take couple of minutes. python print "Supervised learning:" train F.cross entropy print print "Reinforcement learning:" train vanilla pg loss If you compare final accuracy numbers from the both runs the difference is not that much. But if you inspect printed accuracy numbers after each epoch you will notice that with cross-entropy the model converges faster as it gets richer learning signal, the probability of right answer is pushed up during every optimization step, but with policy gradient loss probability of whatever class label we sampled gets pushed up or down depending on the reward. Another important thing to notice here that during evaluation we are not doing any sampling and just taking argmax for both cases to keep the results comparable. In some cases we may want to sample/explore during inference time too, for e.g. to get diverse LLM generations. Next steps Hopefully above code shown the fundamental difference between SL and RL. Once you played around with above approach, I suggest you do the vanilla policy gradient derivation yourself with pen and paper, learn how it uses log-derivative trick to make the whole reward mechanism differentiable — OpenAI's Spinning Up derivation https://spinningup.openai.com/en/latest/spinningup/rl intro3.html deriving-the-simplest-policy-gradient is one good explanation for that these docs are great in general to learn about RL . You can also swap REINFORCE with GRPO based loss function, which will introduce multiple samples and calculation of advantage. There is also this great tutorial paper Invitation to RL https://arxiv.org/pdf/2312.08365 , which is a good followup in the vein of this blog, provides a broader view of RL in a concise and simple manner. I highly recommend this as the next read, and taking it from there.