machine learning
model retraining
new classes
data science
model update

Re-train model with new classes

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Adding new classes is not a minor update. In most classification systems, the output layer, label mapping, metrics, and evaluation set all depend on the class vocabulary. When that vocabulary changes, you usually need more than "one more training run."

The main design question is whether you want a true incremental learner or a standard model that is being expanded to predict additional classes. Most deep-learning classifiers fall into the second category.

What Changes When New Classes Appear

Suppose a model originally predicts two classes, cat and dog. If you now want rabbit as well, several artifacts change at once:

  • the final classifier layer size
  • the label-to-index mapping
  • the training set and validation set
  • any business logic built around class probabilities or thresholds

That is why simply calling fit on new rabbit examples is usually not enough. The old model literally does not have an output neuron for the new class.

The Safe Default: Retrain With Old and New Data

The most reliable approach is to train a new classifier over the full class set using a dataset that includes both old and new classes. This avoids catastrophic forgetting, where the model learns the new class and degrades badly on the old ones.

If you have a useful existing backbone, you do not need to start from random weights. Reuse feature-extracting layers, replace the classifier head, and retrain on the expanded label space.

A Keras Example

This small example trains a two-class model, then replaces the output layer with a three-class head:

python
1import numpy as np
2import keras
3from keras import layers
4
5np.random.seed(0)
6
7x_old = np.random.random((200, 4)).astype("float32")
8y_old = np.random.randint(0, 2, size=(200,))
9
10inputs = keras.Input(shape=(4,))
11x = layers.Dense(16, activation="relu", name="backbone")(inputs)
12outputs = layers.Dense(2, activation="softmax", name="classifier")(x)
13old_model = keras.Model(inputs, outputs)
14
15old_model.compile(
16    optimizer="adam",
17    loss="sparse_categorical_crossentropy",
18    metrics=["accuracy"],
19)
20old_model.fit(x_old, y_old, epochs=3, verbose=0)

Now expand the classifier:

python
1feature_extractor = keras.Model(
2    old_model.input,
3    old_model.get_layer("backbone").output,
4)
5feature_extractor.trainable = False
6
7x_new = np.random.random((300, 4)).astype("float32")
8y_new = np.random.randint(0, 3, size=(300,))
9
10new_inputs = keras.Input(shape=(4,))
11features = feature_extractor(new_inputs)
12new_outputs = layers.Dense(3, activation="softmax", name="new_classifier")(features)
13new_model = keras.Model(new_inputs, new_outputs)
14
15new_model.compile(
16    optimizer="adam",
17    loss="sparse_categorical_crossentropy",
18    metrics=["accuracy"],
19)
20new_model.fit(x_new, y_new, epochs=3, verbose=0)
21print(new_model.predict(x_new[:2], verbose=0).shape)

This is the transfer-learning pattern: keep the learned representation, replace the head, and train on the new class set.

Why Mixing Old and New Examples Matters

If x_new contains mostly rabbit data and very little cat or dog data, the classifier can drift toward the new distribution. Even if the feature extractor is still useful, the new head only learns what you show it.

A better dataset includes:

  • enough examples from old classes to preserve decision boundaries
  • enough examples from the new classes to learn separation
  • validation splits that measure both old-class and new-class performance

When old data is unavailable, class expansion becomes much harder. You may need distillation, rehearsal buffers, or specialized continual-learning methods. Those approaches are real research areas because the problem is genuinely difficult.

Fine-Tuning Versus Full Retraining

Use fine-tuning when:

  • the new classes are related to the old problem
  • the old backbone still extracts useful features
  • you have limited compute or training time

Use full retraining when:

  • the class set changed substantially
  • the data distribution changed, not just the labels
  • the old model was already near its limits

In Keras-style workflows, a common sequence is:

  1. freeze the old backbone
  2. train the new head
  3. unfreeze some or all of the backbone
  4. continue training with a lower learning rate

That matches the standard transfer-learning guidance.

Update the Label Vocabulary Too

Do not forget the metadata. If the original model used:

python
class_names = ["cat", "dog"]

the expanded model must use:

python
class_names = ["cat", "dog", "rabbit"]

This sounds obvious, but many deployment bugs happen here. The numeric model is updated, but the serving code still decodes outputs with the old class list.

What About Incremental Learning?

Some scikit-learn estimators support partial_fit, and some continual-learning methods try to absorb new classes over time. Those are useful when retraining on full historical data is impossible.

But for ordinary neural-network classifiers, class expansion is still usually handled by rebuilding the classifier head and retraining with a mixed dataset. That is the baseline approach you should expect to implement first.

Common Pitfalls

The biggest mistake is training only on the new class and expecting the model to retain old knowledge. That usually causes catastrophic forgetting.

Another is forgetting to resize the final layer. A two-class output cannot produce a third class no matter how much new data you add.

Teams also overlook label metadata, confusion-matrix design, and monitoring dashboards. All of those need the new class list too.

Finally, evaluate per class. A single accuracy number can hide the fact that the new class improved while the old ones collapsed.

Summary

  • Adding new classes usually requires changing the output layer and label mapping.
  • The safest baseline is retraining on a dataset that contains both old and new classes.
  • Reusing the old backbone and replacing the classifier head is a common transfer-learning pattern.
  • Mixing old and new examples reduces catastrophic forgetting.
  • Update deployment metadata and evaluate each class separately after expansion.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.