Getting a foothold in Reinforcement Learning for LLMs — Amit Poonia

(##) 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.<br>```python<br>import torch<br>import torch.nn as nn<br>import torch.nn.functional as F<br>from torch.utils.data import DataLoader<br>from torchvision import datasets, transforms<br>from torch.distributions import Categorical

train_data = datasets.MNIST('./data', train=True, download=True, transform=transforms.ToTensor())<br>test_data = datasets.MNIST('./data', train=False, download=True, transform=transforms.ToTensor())

train_loader = DataLoader(train_data, shuffle=True, batch_size=16)<br>test_loader = DataLoader(test_data, batch_size=100)

class Model(nn.Module):<br>def __init__(self, input_dim: int, hidden_dim: int, output_dim: int) -> None:<br>super().__init__()

self.model = nn.Sequential(<br>nn.Flatten(),<br>nn.Linear(input_dim, hidden_dim),<br>nn.ReLU(),<br>nn.Linear(hidden_dim, output_dim)

def forward(self, input):<br>return self.model(input)<br>```

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.<br>```python<br>@torch.no_grad()<br>def get_accuracy(model: Model, data_loader: DataLoader) -> float:<br>model.eval()<br>correct = 0<br>for x, y in data_loader:<br>correct += (model(x).argmax(dim=1) == y).sum().item()<br>return round(correct / len(data_loader.dataset), 4)

def train(loss_fn, epochs: int=10, learning_rate: float = 2e-4):<br>torch.manual_seed(42)<br>model = Model(784, 100, 10)<br>optimizer = torch.optim.AdamW(params=model.parameters(), lr=learning_rate)<br>for epoch in range(epochs):<br>model.train()<br>for x,y in train_loader:<br>logits = model(x)<br>loss = loss_fn(logits, y)

optimizer.zero_grad()<br>loss.backward()<br>optimizer.step()

train_accuracy = get_accuracy(model, train_loader)<br>test_accuracy = get_accuracy(model, test_loader)

print(f"After epoch {epoch}, train accuracy: {train_accuracy}, test_accuracy: {test_accuracy}")<br>```

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.<br>```python<br>def vanilla_pg_loss(input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:<br>dist = Categorical(logits=input)<br>index_vector = dist.sample()<br>logprob_vector = dist.log_prob(index_vector)<br>reward = torch.where(index_vector == target, 1.0, -1.0)<br>return -(reward * logprob_vector).mean()<br>```

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...