🤖 machine learning

Loss Smoothing for Stable Adaptation Under Distribution Shift

This paper proposes "loss smoothing," a method that interpolates between source and target objectives to enable gradual adaptation under distribution shift, thereby preserving useful features and consistently improving performance across diverse tasks like fine-tuning and reinforcement learning.

Original authors: Darshan Patil, Ekaterina Lobacheva, Razvan Pascanu, Sarath Chandar

Published 2026-07-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Darshan Patil, Ekaterina Lobacheva, Razvan Pascanu, Sarath Chandar

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a master chef who has spent years perfecting a classic French recipe. You know exactly how to handle the ingredients, the timing, and the heat. Now, someone asks you to cook a completely different dish, say, a spicy Thai curry.

The Old Way (The "Hard Switch"):
In traditional machine learning, when it's time to switch tasks, the chef is told to immediately throw away their French cookbook, forget everything they knew about French cooking, and start the Thai recipe from scratch using only the new instructions. They might keep their knife skills (the basic ability to chop), but they are forced to unlearn their French techniques instantly. This sudden change often causes the chef to panic, mess up the new dish, and lose the valuable muscle memory they had built up. In the paper's terms, this is an "abrupt transition" that distorts useful features.

The New Way (Loss Smoothing):
The authors of this paper propose a gentler approach called Loss Smoothing. Instead of throwing away the French cookbook the moment the Thai recipe arrives, they suggest a transition period.

The Blend: At the start of the new task, the chef follows a "blended" instruction. They use 90% of the new Thai instructions and 10% of the old French techniques.
The Fade: As they continue cooking, the recipe slowly shifts. It becomes 80% Thai / 20% French, then 50/50, and finally 100% Thai.
The Result: Because the chef didn't have to make a sudden, jarring jump, they could keep their useful knife skills and general cooking intuition while slowly adapting to the new flavors. The old knowledge acted as a safety net, preventing the chef from making catastrophic mistakes while they learned the new style.

What the Paper Actually Found

The researchers tested this "blending" idea in four different real-world scenarios, acting like a chef trying different new cuisines:

Video Game AI (Reinforcement Learning): Imagine an AI trained to play a game using old, recorded data (offline). When it starts playing live against real opponents (online), it often gets confused. The "Hard Switch" makes it forget the safe strategies it learned from the recordings. Loss Smoothing helped the AI keep those safe strategies while slowly learning to take risks, leading to better scores in games like AntMaze and HalfCheetah.
Language Models (LLMs): Large language models are first trained on a massive mix of internet text (pretraining). Then, they are "fine-tuned" to follow specific instructions. Sometimes, if the model was trained too much on the internet data, it becomes "brittle" and breaks when asked to follow instructions. The paper found that Loss Smoothing (gradually shifting from internet text to instructions) prevented this breakage, allowing the model to stay smart and helpful without losing its general knowledge.
Learning New Tasks (Continual Learning): If a robot learns to recognize cats, then dogs, then birds, a hard switch often makes it forget the cats. Loss Smoothing helped the robot remember the cats while learning the dogs, reducing the "forgetting" that usually happens.
Vision Models: When adapting a model trained on one set of images (like ImageNet) to a new set (like DomainNet), the smooth transition helped the model keep its ability to recognize objects while learning the new style of images.

The Core Takeaway

The paper argues that how you transition from an old task to a new one matters just as much as the new task itself.

The Problem: Jumping straight into a new objective is like slamming on the brakes and then flooring the gas pedal at the same time. It shocks the system.
The Solution: Loss Smoothing is like easing off the gas and gently pressing the new pedal. It keeps the useful parts of the old training alive just long enough to guide the model into the new task, then slowly fades them away so the model can specialize.

The authors conclude that this simple trick—interpolating between the old goal and the new goal—makes models more stable, less likely to "break" during learning, and generally better at adapting to new situations, whether they are robots, language models, or image classifiers.

Important Note: The paper strictly focuses on the mechanics of training these models. It does not claim this method will cure diseases, predict the stock market, or be used in specific future applications beyond the training scenarios they tested (fine-tuning, reinforcement learning, and continual learning). The "magic" is purely in the math of how the learning process is smoothed out.

Technical Summary: Loss Smoothing for Stable Adaptation Under Distribution Shift

Problem Statement

Neural networks are frequently adapted from a source distribution to a target distribution in settings such as fine-tuning, continual learning, and reinforcement learning (RL). Standard adaptation methods typically execute an abrupt "hard switch," discarding the source training objective immediately upon beginning the target phase. The authors argue that this discontinuous transition induces brittle optimization. Specifically, the sudden removal of the source objective can distort learned representations and features that remain useful for the new task, leading to suboptimal performance, catastrophic forgetting, or instability. This issue is particularly pronounced when the target objective is difficult or when the model has undergone "catastrophic overtraining," where parameters become highly sensitive to downstream updates.

Methodology: Loss Smoothing

The paper proposes Loss Smoothing, a lightweight intervention designed to replace the discontinuous switch with a gradual transition between training objectives.

Core Mechanism

Instead of optimizing solely on the target objective $J_t(\theta)$ , the method optimizes a convex interpolation of the source objective $J_s(\theta)$ and the target objective at the start of the adaptation phase:
$J_\alpha(\theta) = (1 - \alpha)J_s(\theta) + \alpha J_t(\theta), \quad \alpha \in [0, 1]$

The interpolation coefficient $\alpha$ is annealed linearly over $\tau$ steps according to the schedule $\alpha(k) = \min(k/\tau, 1)$ .

Early Updates: When $\alpha$ is small, updates remain close to the source training dynamics, preserving useful features and structure from the source distribution.
Later Updates: As $\alpha$ approaches 1, the method recovers the standard target objective, allowing the model to specialize to the new task.

Implementation Details

The method is general and does not assume specific relationships between source and target phases. It requires an estimate of the source objective under current parameters, which can be obtained via:

Replay data: Storing past transitions or inputs (used in RL and continual learning).
Stored targets: Using previously computed value targets or labels.
Direct access: Accessing the previous training distribution (used in LLM fine-tuning).

In cases where the source and target objectives have different scales (e.g., TD3+BC to TD3), the method applies multiplicative annealing to the loss scale to ensure smooth interpolation.

Key Contributions

Formulation of Adaptation: The authors frame adaptation as a transition between training objectives, identifying the "hard switch" to the target objective as a primary source of brittle optimization under distribution shifts.
Loss Smoothing Algorithm: They introduce a simple, path-based interpolation method applicable across input shifts ( $u$ ), learning signal shifts ( $s$ ), and loss family shifts ( $\ell$ ).
Empirical Validation: The method is shown to improve adaptation across diverse settings:
- Continual supervised distribution shifts.
- Pretrained vision model adaptation.
- Offline-to-online and online reinforcement learning.
- Large Language Model (LLM) fine-tuning.
Mechanism Analysis: Evidence suggests that preserving reusable features at the start of adaptation leads to stronger downstream specialization and reduces transient performance collapse (the "stability gap").

Experimental Results

1. Offline-to-Online Reinforcement Learning

Evaluated on D4RL benchmarks (MuJoCo, AntMaze, Adroit), loss smoothing was applied to state-of-the-art offline-to-online methods (TD3+BC, SPOT).

Performance: Loss smoothing achieved the best overall Interquartile Mean (IQM) scores across all suites, outperforming hard-switch baselines and other offline-to-online methods (AWAC, Cal-QL, OPT).
Trade-off: It successfully balanced the need to leverage offline pretraining for exploration (preventing the zero-score failure of training from scratch on difficult tasks like AntMaze) while allowing the policy to move beyond conservative offline constraints as online data became available.

2. LLM Fine-Tuning

Using OLMo-1B models, the authors investigated "catastrophic overtraining," where longer pretraining makes models harder to fine-tune.

Results: Loss smoothing improved the trade-off between in-distribution (ID) instruction-following performance and out-of-distribution (OOD) generalization.
Overtraining Mitigation: In regimes where hard-switch fine-tuning degraded performance due to overtraining, loss smoothing maintained or even increased performance. Notably, the best performance on the Tulu-v1 dataset was achieved using the most "overtrained" checkpoint when smoothed, suggesting the transition path is as critical as the checkpoint itself.

3. Online Reinforcement Learning

Applied to TD3 and PPO algorithms on Brax and DeepMind Control Suite environments.

Stability: Despite existing stabilization mechanisms (replay buffers, target networks), loss smoothing provided significant improvements, with some environments showing nearly double the performance.
Observation: The method helped stabilize the critic adaptation, which is often more difficult than policy adaptation.

4. Continual Supervised Learning & Vision Adaptation

Feature Preservation: In synthetic continual learning tasks (permuted inputs, shuffled labels), loss smoothing preserved task-shared features (e.g., deeper layers in permuted inputs) while adapting task-specific components (e.g., first layers), outperforming EWC and standard replay.
Stability Gap: The method substantially reduced the transient drop in previous-task performance immediately after a task switch.
DomainNet: When adapting an ImageNet-pretrained ViT to DomainNet, loss smoothing achieved a better retention-adaptation frontier than constant replay, indicating that annealing the source objective is more effective than maintaining a fixed replay weight.

Significance and Claims

The paper claims that smoother objective transitions are a broadly useful tool for model adaptation. The central thesis is that brittle adaptation often arises not merely from the difficulty of the target objective, but from the poorly chosen, discontinuous path into it.

Preservation of Structure: By keeping early updates tied to the source objective, loss smoothing prevents the immediate overwriting of useful, task-shared representations.
Path Dependence: The results suggest that the path taken to reach a downstream objective matters. Catastrophic overtraining and feature distortion can be mitigated by how the model transitions from pretraining to fine-tuning, rather than solely by the quality of the pretrained checkpoint.
Generality: Unlike methods that rely on specific data heuristics or architectural constraints, loss smoothing is a general intervention applicable to shifts in inputs, signals, and loss functions.

The authors remain modest regarding limitations, noting that the method requires access to or an estimate of the source objective (replay data, stored targets) and introduces hyperparameters (smoothing horizon) that may need tuning, particularly in online RL where distribution shifts are continuous and less defined. They identify automatic detection of smoothing boundaries and handling incompatible loss scales as directions for future work.