Co-training and catastrophic forgetting

The problem a VLA team meets on day one

A VLA starts as a vision-language model that can describe images, answer questions and follow instructions. To make it drive, or move an arm, you fine-tune it on action data. After a few thousand steps it predicts actions. Ask it a question about an image and it may no longer answer sensibly. The new training overwrote what the old training put there. This is catastrophic forgetting, and it is the reason for a set of recipes in VLA training: mix the old data in, freeze parts of the model, train small adapters.

This module reproduces the effect on a toy small enough to run in seconds, measures how much each remedy keeps, and ends with what to watch for when you apply it to a real model.

Two tasks, one network

Task A and task B each classify inputs into one of six classes, with different random rules (a random two-layer "teacher" for each). The input is 16 random numbers plus a two-valued flag that says which task it is, the way a prompt names a task. A network is trained on A until it is good at it, then adapted to B.

def make_teacher(seed: int) -> Callable[[torch.Tensor], torch.Tensor]:
    """A random two-layer rule that labels inputs with one of CLASSES classes."""
    generator = torch.Generator().manual_seed(seed)
    w1 = torch.randn(DIM, 32, generator=generator) / DIM**0.5
    w2 = torch.randn(32, CLASSES, generator=generator) / 32**0.5
    return lambda x: (torch.tanh(x @ w1) @ w2).argmax(-1)


def make_task(task: int, n: int, seed: int) -> tuple[torch.Tensor, torch.Tensor]:
    """Inputs (n, DIM + 2), with a one-hot task flag appended, and labels (n,)."""
    generator = torch.Generator().manual_seed(seed)
    x = torch.randn(n, DIM, generator=generator)
    flag = torch.zeros(n, 2)
    flag[:, task] = 1.0
    return torch.cat([x, flag], dim=1), make_teacher(task + 1)(x)

A network trained on A scores 82 percent on held-out A inputs (chance is 17 percent). Now fine-tune it on B alone, with the same optimiser and step count:

Task B is learned and task A is gone. Task A's accuracy is not merely lower; it is below chance, because the network now answers confidently with B's rule, and B's rule on A's inputs is unrelated to A's labels. The flag input did not protect task A: the weights that implement A are the same weights that were reshaped to implement B.

Remedy 1: mix old data in (replay and co-training)

If forgetting comes from only seeing B, keep showing A. Replay adds a share of the old task's data to the new training set; co-training is the same idea when the old data is a standing part of the recipe (RT-2, for example, co-fine-tunes its VLM on web data alongside robot data, as the VLA module describes).

In the toy, the share of A's data mixed into B's, and the accuracy it keeps:

Three things to notice. A small share helps a lot: 3 percent, which is 90 examples, moves A from 11 to 39 percent. The recovery is gradual: you pay for each point of A with more data and more training time. And even 30 percent does not reach the from-scratch result, so replay limits forgetting rather than eliminating it. In a real recipe the share, the sampling of the old data (uniform, or favouring examples the model finds hardest) and the order matter, and you must decide what "old data" means when the original pretraining set is not available to you.

Remedy 2: protect the important weights (EWC)

If you cannot keep the old data, you can keep information about it. Elastic weight consolidation (1612.00796) adds a penalty that pulls each weight back towards its old value, scaled by how much that weight mattered for the old task. Importance is estimated by the Fisher information: the average squared gradient of the log-likelihood on the old data.

def fisher_diagonal(
    model: nn.Module, x: torch.Tensor, samples: int = 300
) -> list[torch.Tensor]:
    """Per-weight importance: the mean squared gradient of the log-likelihood.

    A weight whose small change would change the model's predictions on the old task a
    lot has a large value. The label is sampled from the model's own prediction.
    """
    fisher = [torch.zeros_like(p) for p in model.parameters()]
    for i in range(samples):
        model.zero_grad()
        log_probs = nn.functional.log_softmax(model(x[i : i + 1]), dim=-1)
        label = torch.multinomial(log_probs.exp(), 1)[0, 0]
        log_probs[0, label].backward()
        for f, p in zip(fisher, model.parameters(), strict=True):
            f += p.grad**2 / samples
    return fisher


def ewc_penalty(
    model_before: nn.Module, fisher: list[torch.Tensor], strength: float
) -> Callable[[nn.Module], torch.Tensor]:
    """Elastic weight consolidation: strength / 2 * sum fisher * (w - w_old)^2."""
    old = [p.detach().clone() for p in model_before.parameters()]

    def penalty(model: nn.Module) -> torch.Tensor:
        total = torch.zeros(())
        for f, p, o in zip(fisher, model.parameters(), old, strict=True):
            total = total + (f * (p - o) ** 2).sum()
        return 0.5 * strength * total

    return penalty

With the penalty weight $\lambda = 10{,}000$, the network keeps 53 percent on A and reaches 78 percent on B: a clear improvement over fine-tuning, a cost of 3 points on B, and not as good as replay at 10 percent. Smaller $\lambda$ values (10 to 1,000) kept almost nothing in our runs. EWC needs a stored estimate of importance per weight, which is the size of the model, and its quadratic penalty is a coarse picture of the old task. It is a useful idea and a weak default for large models.

Remedy 3: a smaller step

The obvious fix for "the network changed too much" is to change it less. In the toy, a learning rate ten times smaller (3e-4) left task A at 11 percent: the same forgetting. A rate a hundred times smaller (3e-5) kept 55 percent on A, and also learned only 24 percent of B in the same number of steps. A small step size slows both forgetting and learning, and because forgetting is driven by the direction of the updates, not only their size, it often arrives before the new task is learned. In practice a lower rate and a short schedule are cheap and worth doing, and rarely enough.

Remedy 4: freeze the model and train adapters

The most reliable way to keep a capability is not to change the weights that hold it. Low-rank adapters (LoRA, in the efficient fine-tuning module) leave the original weights frozen and learn a small low-rank update on top of each layer. The model with the update switched off is exactly the old model.

class LowRankLinear(nn.Module):
    """A frozen linear layer plus a trainable low-rank update B @ A (zero at the start).

    With ``enabled`` False the layer is exactly the original, so the old task is
    recovered by switching the adapters off.
    """

    def __init__(self, base: nn.Linear, rank: int) -> None:
        """Wraps ``base`` and freezes it."""
        super().__init__()
        self.base = base
        for p in base.parameters():
            p.requires_grad_(False)
        self.down = nn.Parameter(torch.randn(rank, base.in_features) * 0.05)
        self.up = nn.Parameter(torch.zeros(base.out_features, rank))
        self.enabled = True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Base output, plus the low-rank update when enabled."""
        out = self.base(x)
        if self.enabled:
            out = out + x @ self.down.T @ self.up.T
        return out


def with_adapters(model: nn.Module, rank: int) -> nn.Module:
    """A copy of ``model`` whose linear layers are wrapped in low-rank adapters."""
    wrapped = copy.deepcopy(model)
    for i, layer in enumerate(wrapped):
        if isinstance(layer, nn.Linear):
            wrapped[i] = LowRankLinear(layer, rank)
    return wrapped


def set_adapters(model: nn.Module, enabled: bool) -> None:
    """Switches every adapter on or off."""
    for layer in model:
        if isinstance(layer, LowRankLinear):
            layer.enabled = enabled

In the toy, adapters of rank 8 on every layer give 82 percent on A (the unchanged model, measured with the adapters off) and 80 percent on B (with them on), matching training both tasks from scratch. Two honest caveats come with that result.

• You must know which task you are in. Task A is exact because the adapters are off for A and on for B, so something has to switch them. If a single model must serve both with the adapters on, you are back to a changed model, and the forgetting questions return.
• Capacity is bounded by the rank. The adapters learned B here because B resembles A in kind (the same input type, a similar rule). A task very unlike the original may need more rank, or more than adapters.

In a real VLA the closest cousin is freezing the vision encoder or language backbone and training only the action expert, or only adapters. It keeps what the frozen part knows, and limits what the fine-tuning can change.

Measuring forgetting

The toy measures forgetting because it keeps a held-out A set. A real project has to build that set deliberately.

• Keep an "old capabilities" evaluation. For a VLA: image question answering, a few instruction following checks, and the driving or manipulation task it was originally good at. Run it before and after every fine-tune.
• Look at more than the average. Forgetting is uneven; a model can keep its common-sense answers and lose a rare skill. Report per-capability numbers.
• Compare against the mixed baseline. The honest question is not "did it forget less than plain fine-tuning" but "how far is it from training on both together", the last row of the table.
• Watch for forgetting that benchmarks miss. Fluency in a language not covered by the evaluation, safety behaviours, refusal habits and calibration can all move.

Why this matters for a VLA

A driving or robot VLA is the clearest case: its purpose is a new skill (actions), and its value partly comes from what the backbone already knows (language, objects, physical common sense). The recipes above map directly onto training choices you will read about in papers: a data mix that includes general vision-language examples; a frozen or partly frozen backbone; low-rank adapters; and a lower learning rate on pretrained layers. When a paper claims its model keeps its language ability, check whether it evaluated it, and against which baseline.

Reference implementation (PyTorch)

Tests are in libs/tutorial/tests/test_continual_learning.py. They check that a zero-initialised adapter changes nothing before training, that the Fisher estimate is non-negative, that fine-tuning on B destroys A (below 20 percent against a chance level of 17), that more replay keeps more of A, that EWC keeps part of A at some cost to B, that adapters keep A and learn B, and that a smaller step trades forgetting for slower learning. Run uv run --group tutorial python -m libs.tutorial.continual_learning.

Check yourself

1. Why did task A fall below chance rather than to chance after fine-tuning on B?
2. The flag input told the network which task it was. Why did that not protect task A?
3. Replay with 3 percent of A's data moved A from 11 to 39 percent. Why does the first small share help most?
4. What does EWC store, and why does that matter for a large model?
5. A ten times smaller learning rate did not stop forgetting in the toy. A hundred times smaller did, at a price. What was the price?
6. Adapters give an exact old model with the adapters off. What has to be true for that to help in deployment?
7. How would you check that a driving VLA still answers image questions after fine-tuning on trajectories?

Quiz

Answer each question; feedback appears as soon as you choose.

Sources

• Kirkpatrick et al., Overcoming catastrophic forgetting in neural networks, 2016 (EWC).
• Hu et al., LoRA: Low-Rank Adaptation of Large Language Models, 2021.
• Brohan et al., RT-2, 2023, for co-fine-tuning on web and robot data.
• Toy results (82 percent starting accuracy, the replay, EWC, learning-rate and adapter numbers) are from our reference implementation, one seed.