← Back to portfolio
case study

Teaching a Small Language Model to Reason About Unit Conversion

A four-stage exploration of how far you can push a 1.7B-parameter model on a narrow reasoning task — from raw prompting, to in-context chain-of-thought, to LoRA fine-tuning, to a lightweight reinforcement learning loop.

Role
Solo — coursework
Course
Advances in Deep Learning
Stack
Python, PyTorch, HF Transformers, PEFT
CoT Code
01

Overview

For a homework assignment in my Advances in Deep Learning course at UT Austin, I built and compared four different ways of getting a small open-weight language model, SmolLM2 (1.7B parameters), to reliably perform unit conversions (e.g. meters to yards, kilograms to grams). The assignment was structured as a progression: each stage builds on the last, moving from pure prompting techniques toward actual model fine-tuning, ending with a basic reinforcement-learning-style data generation pipeline.

The throughline was measuring how much each additional layer of technique — better prompting, then supervised fine-tuning, then RL-style self-improvement — actually moved the needle on accuracy.

02

The problem

Out of the box, small language models are inconsistent at arithmetic-flavored tasks like unit conversion — they'll often produce a plausible-looking but numerically wrong answer, or fail to isolate a clean final answer from their reasoning at all. The goal was to explore, hands-on, the toolbox modern practitioners use to close that gap: prompt engineering, parameter-efficient fine-tuning, and reinforcement learning from self-generated data — and to understand the tradeoffs (accuracy, compute cost, engineering complexity) each one carries.

03

My approach

I implemented the pipeline in four stages, each with a distinct technical focus:

Stage 1

Batched generation from scratch

Implemented both sequential and batched text generation directly against the model, rather than relying on a high-level pipeline abstraction. Since batching requires all sequences in a batch to share a length, shorter prompts are left-padded and an attention mask tells the model which tokens are real vs. padding. Getting this right unlocked roughly a 8.38x speed improvement over one-at-a-time generation.

Sequential vs. batched generation — 64 prompts
Sequential12.24s · 523 tok/sec
12.24s
Batched (micro-batch 32)1.46s · 4,385 tok/sec
1.46s
8.38x faster with batching, ~8.4x more tokens/sec throughput
Stage 2

In-context chain-of-thought prompting

Built a structured chat prompt — system instructions plus a single well-chosen worked example — that walks the model through reasoning step-by-step before producing a final answer, using HuggingFace's chat template to format everything correctly for SmolLM2's expected input format. No model weights change here; the improvement comes entirely from how the question is framed.

Stage 3

Supervised fine-tuning with LoRA

Rather than fine-tuning the full 1.7B-parameter model, I attached a LoRA (Low-Rank Adaptation) adapter to every linear layer and trained only that — a much smaller set of parameters — to directly predict answers in a constrained <answer>...</answer> format. This keeps the trainable footprint small (the adapter alone, not the full model, is what gets saved and shipped) while still meaningfully shifting the model's behavior on this task.

Stage 4

Rejection-sampling fine-tuning (RFT)

Implemented a simplified version of the RFT algorithm (Yuan et al., 2023): using the chain-of-thought model itself, I generated 16 diverse candidate solutions per question (via higher sampling temperature), kept only the ones that landed on the correct final answer, and used those self-generated, verified reasoning traces as a new training set. This is a lightweight way to get some of the benefit of reinforcement learning — the model learns from its own successful reasoning paths — without implementing a full RL training loop.

04

Challenges

Challenge

Tuning LoRA rank, alpha, learning rate, and training epochs to stay under the file-size budget of 50MB while preserving accuracy.

Challenge

Generating accurate data to increase accuracy of the SFT and RFT models. This required adding more chat templates with diverse examples for the model to reason with.

05

Results

Comparing accuracy across the three learned methods showed a clear, steady climb: each additional technique nearly doubled the gains of the last, with rejection-sampling fine-tuning almost doubling the raw chain-of-thought baseline.

SmolLM2 evaluation accuracy progression across methods
55%
CoT
(in-context)
62%
SFT
(LoRA)
83%
RFT
(rejection sampling)
55%
Chain-of-thought accuracy
62%
SFT (LoRA) accuracy
83%
RFT accuracy — +40pts over prompting alone