From Specialists to a Generalist: Consolidating Knowledge in Neural Networks

When working with neural networks, it is tempting to start with the largest possible model. A sufficiently large network is created, all available data is thrown at it, and the expectation is that it will eventually learn every task that is required.

There is, however, another way of approaching the problem.

Instead of teaching one large network everything at once, many small networks may first be trained for individual situations, potentially even in parallel. One network becomes very good at task A, another one learns task B, and yet another network handles some particularly unpleasant corner case that occurs only occasionally.

At some point one asks obviously:

How do I transfer all of this independently acquired knowledge later into one larger network?

And this leads to a surprisingly interesting way of thinking about how complex neural systems may be built.

• Starting With Specialists
• Learning From Networks Instead of Data
• Why Not Simply Average the Networks?
• The Generalist Needs Context
• A Shared Representation
• Specialists Can Teach More Than Their Original Labels
• Filling the Gaps
• Progressive Consolidation
• Catastrophic Forgetting
• Keeping the Experts Around
• Mixture of Experts Instead of Complete Consolidation
• From Tasks to Goals
• The Missing Ingredient: Time
• From a Bag of Tricks to a World Model
• Most Likely the Giant Network Should Come Last
• References

Starting With Specialists

Assume a number of independently trained models

[ f_1(x), f_2(x), \ldots, f_N(x) ]

have been created. Each of them solves a different problem or covers a different region of the input space. For example $f_1(x)$ might recognize or correct one particular situation, while $f_2(x)$ is responsible for another one.

The straightforward interpretation would be that these networks have to remain separate forever. An external controller could simply select the appropriate network:

[ y = f_k(x) ]

where some external logic decides which $k$ is currently required. This works and is often a very useful architecture. But it is not the only option.

A larger network may instead be trained to reproduce the behaviour of all the smaller ones. The specialists then stop being merely the final solution. They become teachers.

Learning From Networks Instead of Data

Usually a neural network is imagined as being trained from examples

[ (x_i,y_i) ]

An input $x_i$ is presented and the desired output $y_i$ is known. But once a competent network already exists, it can generate these desired outputs itself. For a specialist network $f_k$, arbitrary suitable inputs may be evaluated:

[ y_i = f_k(x_i) ]

This immediately produces a new training dataset

[ D_k = \left\{ (x_i,f_k(x_i)) \right\} ]

If many specialists exist, their datasets may be combined:

[ D = \bigcup_{k=1}^{N}D_k ]

A larger network $F$ can then be trained on this combined knowledge. One possible formulation is

[ F(x,k) \approx f_k(x) ]

where $k$ identifies the task. Training can then be expressed for example with the mean squared error (MSE) as

[ \mathcal L = \sum_{k,i} \left\| F(x_i,k)-f_k(x_i) \right\|^2 ]

Of course the MSE is only an example - for classifiers one often performs distillation on the probabilities, very common with the Kullback–Leibler divergence.

This is one form of knowledge distillation.

The new network does not need to know how the original specialist learned its task, it only requires its behaviour. The knowledge is transferred in function space.

Why Not Simply Average the Networks?

At first sight another idea seems even simpler (and it actually works when models originate from the same initialization or remain in compatible regions of parameter space). If all networks use the same architecture, perhaps their weights could simply be averaged:

[ \theta = \frac{1}{N} \sum_k\theta_k ]

Unfortunately this is usually not equivalent to averaging their knowledge. Neural networks contain a large number of internal symmetries. Consider two consecutive layers

[ W_2W_1 ]

If the hidden neurons are permuted using a permutation matrix $P$, the same function may be written as

[ (W_2P^{-1})(PW_1) ]

The numerical representation has changed substantially while the represented function has not changed at all. Two independently trained networks may therefore implement very similar behaviour using completely different internal arrangements of neurons. Averaging corresponding weights can consequently combine things that have no semantic relationship whatsoever.

The Generalist Needs Context

There is an important difficulty. Assume one expert has learned

[ f_A(x)=+0.2 ]

while another one has learned

[ f_B(x)=-0.2 ]

If the general network only receives $x$, both outputs appear to be correct answers for the same input. A sufficiently confused network may learn

[ F(x)\approx0 ]

This is not a compromise. It is simply wrong for both tasks.

The larger model therefore needs enough information to distinguish the situations. One possibility is explicit task conditioning:

[ F(x,c) ]

where $c$ describes the task or operating mode. The context may contain a categorical task identifier, but it does not have to. It may also describe the surrounding system state:

[ c = (c_1,c_2,\ldots,c_m) ]

This is important because what appears to be the same sensory situation may require entirely different actions depending on the current operating regime. A generalist network can only become general if it is actually given the information required to distinguish those regimes.

A Shared Representation

Once several related tasks are learned by the same network, something more interesting becomes possible. Instead of maintaining a completely separate representation for every problem, a common encoder may be learned:

[ z = E(x) ]

The latent vector $z$ represents relevant information about the current state. Different task-specific heads can then operate on this shared representation:

[ y_k = H_k(z) ]

The architecture becomes conceptually

This is already significantly different from a collection of independent specialist networks. The individual tasks now share a common internal description of the world and can produce useful interactions between them.

Features learned for task $A$ may unexpectedly help task $B$. A property that was originally relevant only to task $C$ may turn out to be useful for task $D$ as well.

The network is no simply storing several independent mappings, it learns structure shared by the tasks.

Specialists Can Teach More Than Their Original Labels

There is another subtle advantage. Suppose a specialist was trained using noisy measurements:

[ y = f^*(x)+\epsilon ]

where $\epsilon$ represents measurement noise. After training, the network may approximate something closer to

[ f(x)\approx E\lbrack y|x \rbrack ]

The trained model has therefore already performed a kind of smoothing. If another network is later trained from the specialist’s output

[ (x,f(x)) ]

it may receive a cleaner (filtered) target than the original raw measurements provided.

This sounds slightly paradoxical. A network is being trained from another network rather than from the supposedly more fundamental experimental data. But the teacher may already contain a useful estimate of the underlying relationship.

Distillation can therefore act not only as compression, but also as a form of knowledge cleanup. Note that this of course only applies to random measurement noise, not to systematic errors.

Filling the Gaps

Specialists introduce another problem. Each expert usually knows only a limited part of the state space. Assume one model was trained for

[ x\in\lbrack 0,1 \rbrack ]

and another one for

[ x\in \lbrack 3,4 \rbrack ]

Neither model says much about

[ x=2 ]

A generalist trained blindly from both experts may produce almost arbitrary behaviour in the gap. This problem becomes especially important in control systems. Interpolation between two valid control regimes is not automatically a valid control regime.

A powerful approach is therefore to deliberately generate states between the known regions. If a simulator exists, this becomes particularly attractive. The training process may explore

[ x\sim p_{\mathrm{exploration}}(x) ]

rather than using only the originally observed states.

The specialists can provide supervision where they are competent, while simulated trajectories, measured data or additional optimization can populate the remaining regions. The resulting system may eventually cover a state space much larger than any individual expert encountered.

Progressive Consolidation

There is no requirement that all specialists have to exist before consolidation begins. Knowledge may instead be accumulated progressively.

First

Later

and later still

The previous generalist itself becomes another teacher. This is a form of continual knowledge consolidation.

The attractive property is that new skills may be added without rebuilding the entire system conceptually from zero. Though there is, besides recursive teacher error accumulation, a serious danger.

Catastrophic Forgetting

If a network is trained only on a new task, parameters useful for older tasks may be overwritten. The model becomes better at the new problem while silently becoming worse at everything it previously knew. This is usually called catastrophic forgetting.

A simple mitigation is replay. Representative states from previously learned tasks are retained in a replay dataset $D_{\mathrm{replay}}$. Training then uses both $D_{\mathrm{new}}$ and $D_{\mathrm{replay}}$.

One possible loss is

[ \mathcal L = \lambda_{\mathrm{new}}\mathcal L_{\mathrm{new}} + \lambda_{\mathrm{old}}\mathcal L_{\mathrm{distill}} ]

The first term teaches the new behaviour, the second one penalizes changes to behaviour that should remain intact. One can think of this as periodically reminding the network to learn the new thing but retain it’s old capabilities. That is surprisingly close to the problem encountered in many continually learning systems.

Keeping the Experts Around

After consolidation, the specialist networks do not necessarily have to be deleted. Keeping them can be extremely useful. They can act as regression references.

For a set of validation states $x_i$, the consolidated model may periodically be compared against the old expert:

[ \Delta_i = \left\| F(x_i,k)-f_k(x_i) \right\| ]

If $\Delta_i > \epsilon$ something has changed significantly. This gives the system a permanent collection of known-good behaviours against which future versions can be tested.

Mixture of Experts Instead of Complete Consolidation

There is also no requirement that everything has to be compressed into one homogeneous network. Another possibility is a learned router.

The system becomes

The router learns which expert or combination of experts should process a given situation. This is broadly the idea behind a Mixture of Experts system.

It offers an interesting compromise. Shared components can learn common representations while specialist components retain their individual competence. Instead of asking whether the final system should consist of specialists or a generalist, both can coexist.

This may actually be the better architecture for sufficiently complicated physical systems.

From Tasks to Goals

The most fascinating transition occurs when the network is no longer explicitly told which task to execute. Initially the system may be formulated as

[ F(x,\text{task})\rightarrow y ]

But tasks are only human descriptions of desired transformations. A more general formulation is

[ F(x,g)\rightarrow y ]

where $g$ describes a goal. Instead of saying “execute procedure C” one could eventually specify “move the system toward a state with these properties”.

A global objective may, for example, be written as

[ J(x) = \alpha J_1(x) +\beta J_2(x) -\gamma J_3(x) ]

The specialists then become demonstrations of trajectories that improve different components of $J$. At this point the architecture starts looking substantially less like a collection of classifiers and substantially more like a controller.

The Missing Ingredient: Time

For many physical systems, the present observation does not fully describe the system. The same measured state can behave differently depending on how it was reached. Hysteresis exists, components settle over time, temperatures drift, mechanical systems exhibit backlash, charges accumulate, there are delays, etc.

Consequently $a_t = \pi(x_t)$ may fundamentally be insufficient. The appropriate action can depend on history:

[ a_t = \pi(x_t,x_{t-1},\ldots,a_{t-1},a_{t-2},\ldots) ]

This history can be compressed into an internal state

[ h_t = F(h_{t-1},x_t,a_{t-1}) ]

and the controller becomes

[ a_t = \pi(h_t,g) ]

The internal state $h_t$ can be interpreted as the networks current belief about the physical system. At this point something rather interesting has happened. The network is no longer learning individual tasks, it is beginning to construct a model of the system it controls.

From a Bag of Tricks to a World Model

This suggests a useful progression.

A project may begin with many small networks (specialist). These networks solve narrow problems very well. Their behaviour is then consolidated:

Additional specialists are added over time:

Eventually temporal state and goals may be included:

And somewhere along this path the nature of the system changes. What originally consisted of many independent tricks starts acquiring a common representation of the environment.

The specialists knew how to react, the consolidated system may eventually begin learning what kind of situation it is in. That is a considerably more powerful abstraction.

Most Likely the Giant Network Should Come Last

There is a funny reversal in all of this. Modern machine learning often encourages starting with the giant model.

• Collect everything
• Train everything jointly
• Scale until the problem disappears

But for engineering systems it may sometimes be much more practical to do almost the opposite.

• First solve narrow problems
• Make the individual pieces work
• Understand the strange corner cases
• Train small experts until each one is reliable
• Only then ask how their knowledge can be combined

The giant network is not the starting point, it is the consequence of accumulated working solutions. Instead of hoping that one enormous network somehow discovers every useful behaviour at once, the knowledge can grow piece by piece:

[ \text{experience} \rightarrow \text{specialization} \rightarrow \text{distillation} \rightarrow \text{shared representation} \rightarrow \text{generalization}. ]

The small networks are therefore not disposable prototypes on the way toward the real model, they can be exactly how the real model learns what it should become.

References

• Goodfellow, Bengio, Courville, Deep Learning (Adaptive Computation and Machine Learning series)
• Sutton, Barto, Reinforcement Learning, second edition: An Introduction (Adaptive Computation and Machine Learning series)
• Huyen, Designing Machine Learning Systems: An Iterative Process for Production-Ready Applications

This article is tagged: Artificial Intelligence, Tutorial, How stuff works, Large Language Models, Machine learning, Computational linear algebra, LLM, Optimization