Training a Small Language Model: What I Learned About Probability, Context and Bias

Training a Small Language Model changed the way I think about AI.

Not because I suddenly learned that models predict the next token. I already knew that at a conceptual level. The interesting part was understanding what sits underneath that seemingly simple statement.

Probability. Context. Data representation. Repetition. Error. Adjustment.

And ultimately, judgement.

The deeper I went, the clearer something became: a language model does not need to contain a perfect description of reality to be useful. In many cases, attempting to encode every exception would actually make the system worse.

Consider a simple statement:

“Most birds fly.”

It is not universally true. Penguins do not fly. Ostriches do not fly. Very young birds may not yet fly. An injured bird may temporarily be unable to fly.

We could attempt to create an enormous deterministic rule describing every exception.

Or we could represent something much closer to how humans actually reason:

Given what we know, a bird is highly likely to fly.

That shift from absolute rules to probabilistic reasoning is foundational to machine learning.

Language is probability

British linguist John Rupert Firth famously said:

You shall know a word by the company it keeps.

That line might as well be an introduction to modern language modelling.

Words are not meaningful merely because they exist in a vocabulary. Their meaning emerges from context.

Consider:

“Apple released a new…”

The probabilities for the next word might strongly favour words such as:

iPhone

MacBook

chip

device

Now consider:

“She sliced the apple with a…”

The probability distribution changes completely.

The word “apple” is identical.

The context around it is not.

This is one of the most important things I have learned while studying language models:

Contextual cues shape probability distributions.

Almost everything else follows from there.

Starting with N-grams

One of the simplest ways to understand language modelling is through N-grams.

A unigram considers one token.

A bigram considers sequences of two tokens.

A trigram considers sequences of three tokens.

The effective context window is therefore always:

n - 1

A bigram uses one previous token as context.

A trigram uses two.

If our dataset frequently contains:

“a staple food”

A trigram model can estimate:

P(food | a staple)

We count how frequently the sequence “a staple food” appears relative to other sequences beginning with “a staple”.

This is already a language model.

Given a context, assign probabilities to possible continuations.

The problem is that N-grams run into a brutal limitation: data sparsity.

As the context becomes longer, the number of possible combinations explodes.

A sentence may be perfectly reasonable while never having appeared in the training corpus before. An N-gram model consequently has little or no information with which to estimate its probability.

Transformers attack this problem very differently.

They do not simply memorise fixed sequences.

They learn representations.

First, language becomes numbers

Models cannot process words directly.

The text must first become tokens, and the tokens must become numbers.

A simple tokenizer might take:

“Abeni is brilliant”

and split it into:

["Abeni", "is", "brilliant"]

A vocabulary is then created from the unique tokens appearing throughout the dataset.

return sorted(list(set(tokens)))

Each token receives an integer ID.

Conceptually:

Abeni      -> 0
brilliant  -> 1
is         -> 2

The tokenizer can now encode text:

"Abeni is brilliant"

into something resembling:

[0, 2, 1]

and decode those numbers back into language.

This sounds almost trivial, but it represents an important transformation.

Language has become something mathematics can operate on.

Preparing sequences for training

Neural networks generally expect structured numerical inputs.

Suppose our encoded paragraph is:

[14, 52, 91, 8, 27]

For next-token prediction, we shift the sequence.

The model input becomes:

[14, 52, 91, 8]

The target becomes:

[52, 91, 8, 27]

In code:

input_sequences = padded_sequences[:, :-1]
target_sequences = padded_sequences[:, 1:]

The model is effectively being taught:

Given token 14, predict 52.

Given 14 and 52, predict 91.

Given 14, 52 and 91, predict 8.

And so on.

Across enough examples, the model begins learning the statistical structure of language.

The problem of different sequence lengths

Real text is messy.

One paragraph may contain eight tokens.

Another may contain eighty.

But training data is typically organised into matrices where each row represents an example. Those rows therefore need compatible dimensions.

There are two obvious solutions.

We can truncate long sequences.

Or we can pad short ones using a special <PAD> token.

Neither extreme is ideal.

Truncating every paragraph to the length of the shortest paragraph could throw away enormous amounts of information.

Padding everything to an absurdly long sequence wastes computation.

The practical solution is usually a combination.

Choose a target sequence length.

Truncate sequences beyond it.

Pad sequences below it.

This small preprocessing decision is a good example of something that appears repeatedly in machine learning: elegant systems are often built from pragmatic compromises.

From sequences to batches

Training examples are normally grouped into batches.

For example:

batch_size = 32

Instead of processing the entire dataset simultaneously, the model processes thirty-two examples, calculates its error, adjusts its parameters and moves to the next batch.

Before this happens, we commonly shuffle the dataset.

tf_dataset = tf_dataset.shuffle(buffer_size=len(input_sequences))

This prevents the model from repeatedly seeing examples in exactly the same order.

Once batched, the dataset becomes a collection of tensors — matrices containing token IDs for groups of paragraphs.

Now training can begin.

What is actually being trained?

This is where the process becomes particularly interesting.

At the beginning, the model's parameters are largely random.

It knows nothing.

It receives an input sequence and produces probabilities for what it believes should come next.

At first, those predictions are terrible.

The training process measures how wrong the model was using a loss function.

Optimisation then changes the parameters slightly so that future predictions should become better.

The next batch arrives.

Predict.

Measure error.

Adjust.

Repeat.

After the model has processed the entire training dataset once, it has completed an epoch.

Then we do it again.

And again.

Sometimes hundreds of times.

Training a neural network can therefore be reduced to a surprisingly simple conceptual loop:

Predict.

Measure.

Correct.

Repeat.

The complexity lies in what is being adjusted and how those adjustments interact across millions or billions of parameters.

Learning rate matters

One important variable is the learning rate.

If the learning rate is too high, the model can make enormous corrections and overshoot useful solutions.

If it is too low, learning can become painfully slow.

Imagine trying to reach the lowest point in a valley while blindfolded.

Take enormous jumps and you may repeatedly leap past the bottom.

Take microscopic steps and you will eventually arrive, but perhaps sometime next century.

The learning rate determines the size of those steps.

A value such as:

learning_rate=1e-4

controls how aggressively the model updates itself during training.

The mathematics becomes more sophisticated, but the intuition remains simple.

Watching a model learn is fascinating

One particularly satisfying technique is generating text periodically during training.

For example, start with:

Abeni,

and ask the model to generate ten tokens every ten training iterations.

Initially, the results may resemble nonsense.

Later, sentence fragments emerge.

Eventually, structure appears.

This is not the model retrieving a stored sentence.

It is repeatedly calculating probability distributions over possible next tokens.

That distinction matters.

A model does not predict one answer

A language model does not ordinarily think:

“The next word is definitely X.”

It produces something closer to:

token A: 0.42
token B: 0.21
token C: 0.14
token D: 0.06
...

This is a probability distribution.

Generation then requires selecting from that distribution.

One approach is greedy decoding.

Always choose the token with the highest probability.

This produces predictable output, but can become repetitive.

Alternatively, we can sample from the distribution.

A token with 20% probability can occasionally beat one with 50%.

That introduces variation.

Suddenly we encounter another important balancing act:

Determinism versus randomness.

Accuracy versus creativity.

Predictability versus diversity.

Modern generative systems are constantly managing some version of this tension.

Context is everything

Longer context generally provides more information from which the model can make its next prediction.

Compare:

bank

with:

He deposited money at the bank

and:

They sat beside the river bank

The additional context dramatically changes the likely continuation.

This is why larger context windows can be so useful.

But longer context is not free.

More information also means more computation, more relationships to evaluate and potentially more irrelevant information.

Again, AI engineering becomes a question of balance.

More context can improve prediction.

More context can also increase complexity.

Information is useful only when the system can interpret it properly.

Models are compressed representations of patterns

Perhaps the most important conceptual shift for me was recognising that models are fundamentally representations of patterns.

They absorb information from data and encode relationships within their parameters.

Language patterns.

Semantic associations.

Structural relationships.

Statistical regularities.

And unfortunately, human biases.

The model learns what exists in its training data, not necessarily what ought to exist in society.

That distinction is crucial.

Bias is not an edge case

Suppose a training dataset contains many more references to female nurses than male nurses.

And many more references to male doctors than female doctors.

A model trained on this text may assign higher probabilities to:

“the nurse… she”

and:

“the doctor… he”

The model has discovered a statistical pattern.

Mathematically, it may even be modelling its dataset correctly.

Socially, blindly reproducing that pattern would reinforce a stereotype.

That exposes a fundamental challenge in AI.

The most statistically probable response is not automatically the most appropriate response.

We therefore cannot treat model training as simply maximising predictive accuracy.

AI systems interact with humans.

That introduces values, expectations and consequences.

Some patterns should not be amplified

Consider a harmless probabilistic observation:

“Children often avoid green vegetables.”

There may be data supporting it.

But an educational AI repeatedly telling children that avoiding vegetables is normal or desirable would be counterproductive.

Similarly, datasets inevitably contain stereotypes about professions, nationalities, communities, genders and social groups.

A sufficiently capable model will detect them.

The engineering challenge is not preventing the model from discovering patterns.

It is preventing statistical correlations from automatically becoming normative judgements.

That distinction between:

what is statistically present

and

what the system should endorse

is one of the most important areas of modern AI development.

More information does not automatically solve ethics

A classic example involves autonomous vehicles.

If a vehicle must make an unavoidable decision involving human lives, could more contextual information help it make the “correct” decision?

At first this seems appealing.

The system could theoretically consider age, circumstances, dependencies and consequences.

But there is a dangerous boundary here.

Allowing an algorithm to calculate one human's societal worth against another based on education, criminal history, employment, family background or perceived contribution would encode extraordinary moral and social assumptions into software.

More data does not make those assumptions objective.

An AI system could become more informed while simultaneously becoming less just.

For safety-critical systems, the better principle is therefore not:

collect enough information to decide whose life matters more.

It is:

design the system to minimise harm without assigning different intrinsic values to human lives.

This is an important correction to an intuition I initially found attractive while learning these concepts.

Information is power.

But information does not remove the need for principles governing how that power is used.

Language itself can carry culture

Another fascinating challenge is that language is not culturally neutral.

Research has shown that models can respond differently to ethical questions depending on the language in which those questions are asked.

That makes sense when you think about the training process.

Different languages contain different literatures, histories, norms, cultural assumptions and distributions of ideas.

A multilingual model is therefore not merely translating words.

It is navigating overlapping statistical representations of human cultures.

That makes alignment dramatically more difficult than simply adding a list of prohibited outputs.

Base models and instruction models

Another distinction that became much clearer to me during this process is the difference between a base model and an instruction-tuned model.

A useful simplification is:

Base model = knows language.

Instruction model = knows language and has been trained to respond usefully to people.

A base model primarily learns continuation.

Give it:

The capital of France is

and it learns that “Paris” is an extremely likely continuation.

But conversational assistants require additional behavioural training.

Answer the question.

Follow instructions.

Respect formatting.

Avoid harmful behaviour.

Recognise when uncertainty matters.

That additional training transforms a language predictor into something resembling an assistant.

The frameworks are different layers of the same ecosystem

Working through model training also helped clarify the role of the major machine-learning frameworks.

Keras provides a high-level deep-learning interface that is excellent for learning, experimentation and clean model definitions.

PyTorch offers tremendous flexibility and dominates large areas of AI research and custom model development.

TensorFlow remains a broad machine-learning ecosystem with mature production and deployment infrastructure.

JAX provides extremely high-performance numerical computing and automatic differentiation and has become important in large-scale model research.

Scikit-learn remains extraordinarily useful for classical machine-learning problems such as regression, decision trees, clustering and preprocessing.

Hugging Face has effectively become an ecosystem around pretrained transformers, datasets, tokenizers and model fine-tuning.

And with Keras 3, even some of these boundaries are becoming less rigid because Keras can operate across JAX, PyTorch and TensorFlow backends.

The ecosystem looks fragmented from the outside.

Once you understand the layers, it begins making much more sense.

What training a small model teaches you

Training an SLM is valuable even if you never intend to build the next frontier model.

Because suddenly many concepts stop being buzzwords.

Tokens are no longer abstract.

Probability distributions are no longer abstract.

Context windows are no longer abstract.

Batch size, epochs, learning rates, padding, loss and sampling all become pieces of one connected system.

You begin seeing the machinery underneath the interface.

And perhaps most importantly, you realise that the intelligence we experience from these systems emerges from a deceptively simple objective repeated at extraordinary scale:

Given everything I have seen so far, what is most likely to come next?

Then predict.

Measure.

Adjust.

Repeat.

My biggest takeaway

I started this learning journey thinking primarily about models.

I have increasingly found myself thinking about information.

A model's ability to generate coherent output depends heavily on the quality and relevance of the context available to it.

More relevant information can improve prediction.

Better representations can reveal deeper patterns.

Better training can refine those representations.

But information alone is not intelligence.

Probability alone is not judgement.

And statistical correctness alone is not wisdom.

Training an SLM makes that distinction surprisingly tangible.

We are building machines capable of modelling increasingly large portions of human knowledge, language and behaviour.

The difficult part will not merely be teaching them how to predict.

It will be deciding which predictions should become actions, which patterns should be challenged, how uncertainty should be represented and where human principles must override statistical probability.

That, to me, is where the truly interesting engineering begins.