Building a Personalized AI Text Prediction System with LSTM and GRU

What "text prediction" actually is

The keyboard-suggestion feature on a phone, or the "type ahead" in an email client, is solving one specific problem: given the sequence of words (or characters) so far, predict a probability distribution over what comes next. It's a sequence modeling problem, and recurrent neural networks — specifically LSTM and GRU — were the standard tool for it for years before attention-based models took over the largest-scale version of this task. For a personalized, resource-constrained predictor (running on-device, or trained on a small, specific corpus), they're still a genuinely good, cheap, well-understood choice.

Why not a plain RNN

A vanilla RNN updates a hidden state at every step by combining it with the current input through a single transformation. In theory it can remember arbitrarily far back. In practice, gradients passed back through many time steps during training either vanish toward zero or explode, which means a plain RNN is bad at learning dependencies more than a handful of tokens apart — exactly the kind of dependency you need for "the subject of this sentence was plural, so the verb three words later should be too."

LSTM and GRU both exist to fix that, via gating: instead of overwriting the hidden state wholesale at every step, the network learns how much of the past to keep and how much of the new input to let in.

LSTM: three gates and a separate cell state

An LSTM cell keeps two things across time steps — a hidden state and a cell state — and uses three gates to control them:

• Forget gate — how much of the existing cell state to keep
• Input gate — how much of the new candidate information to add
• Output gate — how much of the cell state to expose as the hidden state

import torch.nn as nn

class LSTMPredictor(nn.Module):
    def __init__(self, vocab_size, embed_dim=128, hidden_dim=256, num_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(
            embed_dim, hidden_dim, num_layers,
            batch_first=True, dropout=0.2,
        )
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x, hidden=None):
        embedded = self.embedding(x)
        output, hidden = self.lstm(embedded, hidden)
        logits = self.fc(output)
        return logits, hidden

The separate cell state is what gives LSTM its long-range memory — it's a slower-changing "conveyor belt" of information that the gates read from and write to, rather than something that gets fully recomputed every step.

GRU: two gates, no separate cell state

A GRU merges the cell state into the hidden state and reduces three gates to two:

• Reset gate — how much of the previous hidden state to ignore when computing the new candidate
• Update gate — how much to blend the previous hidden state with the new candidate

class GRUPredictor(nn.Module):
    def __init__(self, vocab_size, embed_dim=128, hidden_dim=256, num_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.gru = nn.GRU(
            embed_dim, hidden_dim, num_layers,
            batch_first=True, dropout=0.2,
        )
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x, hidden=None):
        embedded = self.embedding(x)
        output, hidden = self.gru(embedded, hidden)
        logits = self.fc(output)
        return logits, hidden

Fewer parameters per cell, fewer matrix multiplications per step. The practical tradeoff is well established in the literature and worth stating plainly instead of pretending there's a universal winner: GRU tends to train faster and needs less data to reach a reasonable result, which makes it the more forgiving choice for smaller or personal datasets. LSTM's extra capacity tends to pay off more clearly on longer sequences and larger corpora, where the separate cell state has more to actually remember. If you're not sure which fits your dataset, that asymmetry — GRU as the safer default for small data, LSTM worth trying once you have more of it — is a reasonable way to decide instead of guessing.

That said, "tends to" is doing real work in that sentence — see the actual measured results below, where GRU didn't train faster at all.

Training loop essentials

The architecture is the easy 20%. Getting a next-token predictor to train well is mostly about a few unglamorous details:

def train_epoch(model, dataloader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        logits, _ = model(inputs)
        loss = criterion(logits.view(-1, logits.size(-1)), targets.view(-1))
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(dataloader)

Two details that matter more than they look: gradient clipping (clip_grad_norm_), because exploding gradients are still a real risk even with gated architectures, especially early in training; and teacher forcing during training (feeding the true previous token as input, not the model's own prediction), which stabilizes training but needs to be gradually reduced or dropped at inference time — otherwise the model has never practiced recovering from its own mistakes, and error compounds token by token during actual generation.

For evaluation, perplexity (the exponential of the average cross-entropy loss) is the standard metric — lower is better, and it's directly comparable across model configurations trained on the same data, which makes it far more useful during development than eyeballing sample generations.

What "personalized" changes

A generic text predictor trained on a huge public corpus learns the statistics of language in general. A personalized one has to learn a specific person's vocabulary, phrasing habits, and typical sentence lengths — from a dataset that is, by construction, much smaller. That changes the practical decisions:

• Fine-tune, don't train from scratch. Start from a model pretrained on a large general corpus and continue training on the personal dataset with a lower learning rate. Training a good language model from scratch on one person's messages is not realistic with the amount of data available.
• Regularize harder. A small, personal dataset overfits fast. More dropout, fewer parameters, and early stopping based on a held-out validation slice matter more here than in the large-corpus case.
• Treat the data as sensitive by default. Personal writing is about as identifying as data gets. Train and serve on-device where possible, encrypt anything stored server-side, and don't fold a user's personal fine-tuning data into any shared or future training set without clear, explicit consent — this is a case where the privacy design isn't a compliance checkbox, it's central to whether the product should exist in the current form at all.

Real Results From an Actual Implementation

Everything above is the theory, and theory is where this article originally stopped — I hadn't trained the comparison myself yet. I have now: a full LSTM/GRU next-word prediction pipeline (Keras/TensorFlow, not the PyTorch sketched above — the architecture is the same idea, different framework), trained on the DailyDialog conversational dataset, with both models run to completion under identical conditions so the numbers are actually comparable.

GRU came out slightly ahead on both accuracy and perplexity — and took longer to train, not shorter. That directly contradicts the "GRU trains faster" heuristic stated earlier as a general tendency, on this exact dataset and setup. Neither result invalidates the other: the heuristic is still a reasonable prior for picking a starting point, but this run is the concrete reminder of why "measure it yourself" was the actual advice, not "GRU wins by default." Full code, the trained models, and the Streamlit comparison interface are on GitHub, and the project is also written up as a case study with more detail on the pipeline.

The honest summary

LSTM and GRU aren't obsolete for this — for a small, personalized, resource-constrained predictor, they're often the more appropriate choice than reaching for a large transformer, precisely because they need less data and less compute to reach a usable result. The architecture is genuinely the easy part. The dataset size, the fine-tuning strategy, and the privacy handling are where a "personalized" version of this system is actually won or lost — and where the numbers don't match the textbook heuristic, that's worth trusting the numbers.