LLM Training Tutorial

An educational project that solves the same regression task — predicting medical insurance charges — across six ML frameworks, so you can compare their APIs and trade-offs side by side.

Examples

File	Framework	Model type	GPU
sklearn_regression.py	scikit-learn	MLP / Linear Regression	CPU only
pytorch_regression.py	PyTorch	MLP / Linear	CUDA auto-detected
tensorflow_regression.py	TensorFlow / Keras	MLP / Linear	CUDA auto-detected
xgboost_regression.py	XGBoost	Gradient boosted trees / gblinear	CUDA auto-detected
lightgbm_regression.py	LightGBM	Gradient boosted trees / linear tree	CPU only (GPU needs custom build)
jax_regression.py	JAX + Flax	MLP / Linear	CUDA (Linux/WSL2 only)

All examples share the same data pipeline via helpers.py: CSV loading, 80/20 split, StandardScaler feature normalisation, and log-transform on the target.

Benchmark results

SAVE_MODEL=false for all runs. Ranked by MSE — lower is better.

MODEL_TYPE=mlp

Rank	Framework	Device	Test R²	Train R²	MSE (×10⁸)
1	XGBoost	CUDA	0.861	0.923	0.1852
2	LightGBM	CPU	0.862	0.882	0.1914
3	sklearn	CPU	0.842	0.850	0.2394
4	JAX	CPU	0.849	0.851	0.2524
5	PyTorch	CUDA	0.839	0.841	0.2586

MODEL_TYPE=linear

Rank	Framework	Device	Test R²	Train R²	MSE (×10⁸)
1	LightGBM	CPU	0.857	0.887	0.2128
2	sklearn	CPU	0.784	0.763	0.6981
2	XGBoost	CUDA	0.784	0.763	0.6981
2	PyTorch	CUDA	0.784	0.763	0.6980
2	JAX	CPU	0.784	0.763	0.6981

LightGBM ranks first in linear mode because its linear_tree=True fits a linear regression at each leaf of every boosted tree — an ensemble of linear models — rather than a single global linear function. The other four frameworks all solve identical linear regression and converge to the same result.

Reading the output

Each script prints a summary after training. Here is what every term means.

Loss — the training error the model is minimising, reported every 50 epochs (neural network scripts only). Computed as Mean Squared Error in log-space: the average squared difference between the model's predictions and the true values, both expressed as log(1 + charges). Lower is better. It will never reach zero on real data because no model fits perfectly.

LR (learning rate) — controls how large a step the optimiser takes each epoch when adjusting the model's weights. A high LR moves fast but can overshoot; a low LR is precise but slow. PyTorch uses cosine annealing, which smoothly decays LR from its starting value to 0 over all epochs.

Test R² — how well the model generalises to data it has never seen (the held-out 20% test split). R² = 1 means perfect predictions; R² = 0 means the model does no better than always predicting the average charge; negative R² means it is actively worse than that baseline. This is the primary quality metric.

Train R² — the same score but measured on the training data the model learned from. If Train R² is much higher than Test R², the model has memorised the training set without learning the underlying pattern (overfitting). Ideally the two values are close.

MSE (×10⁸) — Mean Squared Error in original dollar units, divided by 10⁸ to keep the number readable. It answers: on average, how many dollars² off is each prediction? Unlike R², MSE is an absolute number so you can compare it across datasets. Lower is better.

Predicted / Actual — a single spot-check on the first row of the test set. The model outputs a predicted insurance charge; the actual charge for that patient is shown alongside it. Because the model trains on log(1 + charges), the prediction is converted back to dollars with expm1 before printing. This check catches obvious failures (e.g. a model predicting negative charges) that aggregate metrics can hide.

Shared pipeline (helpers.py)

Function	Purpose
load_and_preprocess(path)	Load CSV, encode categoricals, separate features/target
split_data(X, y, dtype)	Train/test split + log1p target transform
fit_scaler(X_train, X_test, dtype)	Fit StandardScaler on train, transform both splits
restore_scaler(mean, scale)	Reconstruct a scaler from saved arrays (PyTorch / JAX checkpoints)
load_env()	Read SAVE_MODEL and MODEL_TYPE from .env
get_cuda_device()	Detect CUDA availability via PyTorch

Dataset

insurance.csv — 1,339 rows of patient medical/demographic data. Columns: age, sex, bmi, children, smoker, region, charges

topical_chat.csv — 188,378 rows of multi-turn conversations with sentiment labels, reserved for future NLP examples.

Setup

Requires Python 3.10–3.13.

# Install uv if you don't have it
pip install uv

# Create virtual environment and install all dependencies
python -m uv sync

# Copy the environment template and edit as needed
cp .env.template .env

Running examples

python -m uv run sklearn_regression.py
python -m uv run pytorch_regression.py
python -m uv run xgboost_regression.py
python -m uv run lightgbm_regression.py
python -m uv run jax_regression.py

TensorFlow has no wheels for Python 3.13 or 3.14. tensorflow_regression.py is included as a readable reference only. Running it requires a Python ≤ 3.12 environment.

Testing saved models

test_model.py loads a saved checkpoint, runs two hardcoded patients through it, and asserts each predicted charge falls within a known plausible range.

Train first with SAVE_MODEL=true, then set TEST_FRAMEWORK in .env and run:

SAVE_MODEL=true python -m uv run sklearn_regression.py
python -m uv run test_model.py

TEST_FRAMEWORK	Checkpoint files required
sklearn	sklearn_model.joblib
pytorch	pytorch_model.pt
xgboost	xgboost_model.ubj, xgboost_scaler.joblib
lightgbm	lightgbm_model.joblib
jax	jax_model.joblib

The two test patients are hardcoded in the script:

Patient	Profile	Expected charge
A	Age 25, female, non-smoker, BMI 22, northeast	$1,500–$7,000
B	Age 60, male, smoker, BMI 36, southeast	$25,000–$65,000

Environment variables (.env)

Variable	Values	Effect
SAVE_MODEL	true / false	Save trained model to disk and reload it on the next run instead of retraining. Default: false.
MODEL_TYPE	mlp / linear	Switch between the complex model (neural network / boosted trees) and a simple linear model. Default: mlp.

GPU / CUDA notes

PyTorch — auto-detects CUDA 12.4. Prints Using device: cuda or cpu at startup.

XGBoost — auto-detects CUDA via device="cuda". Falls back to CPU silently.

TensorFlow — auto-detects GPU. Windows native GPU support limited to TF ≤ 2.10; use WSL2 for newer versions.

JAX — CPU-only by default. For GPU install jax[cuda12] on Linux/WSL2.

LightGBM — CPU only from PyPI. GPU requires building from source with CUDA enabled.

scikit-learn — CPU only; no GPU support.

Training data

Came from (Medium)[https://python.plainenglish.io/linear-regression-for-predicting-medical-patients-insurance-expenses-with-scikit-learn-ecb63527e524] and more specifically (Kaggle)[https://www.kaggle.com/datasets/mirichoi0218/insurance]