arXiv:2605.12419v1 [cs.CL] 12 May 2026
O RBIT: Preserving Foundational Language Capabilities in GenRetrieval via Origin-Regulated Merging
Neha Verma∗ Johns Hopkins University [email protected]
Nikhil Mehta∗ Google DeepMind [email protected]
Shao-Chuan Wang Google DeepMind
Naijing Zhang Google
Alicia Tsai Google DeepMind
Lukasz Heldt Google
Lichan Hong Google DeepMind
Aniruddh Nath Google
Ed Chi Google DeepMind
Li Wei Google
Xinyang Yi Google DeepMind
Abstract Despite the rapid advancements in large language model (LLM) development, fine-tuning them for specific tasks often results in the catastrophic forgetting of their general, language-based reasoning abilities. This work investigates and addresses this challenge in the context of the Generative Retrieval (GenRetrieval) task. During GenRetrieval fine-tuning, we find this forgetting occurs rapidly and correlates with the distance between the fine-tuned and original model parameters. Given these observations, we propose O RBIT, a novel approach that actively tracks the distance between fine-tuned and initial model weights, and uses a weight averaging strategy to constrain model drift during GenRetrieval fine-tuning when this inter-model distance exceeds a maximum threshold. Our results show that O RBIT retains substantial text and retrieval performance by outperforming both common continual learning baselines and related regularization methods that also employ weight averaging.
1
Introduction
The Generative Retrieval (GenRetrieval) paradigm (Rajput et al., 2023; Tay et al., 2022) has demonstrated considerable efficacy in sequential recommendation tasks. GenRetrieval introduces items or queries as tokenized ID sequences enabling an autoregressive model to sequentially predict relevant items based on preceding context. However, fine-tuning large language models (LLMs) for this specialized purpose introduces a critical challenge known as catastrophic forgetting. This phenomenon, where a model loses previously learned information after training on a new task, causes a significant degradation of the LLM’s pre-existing general-purpose capabilities. Such a trade-off between desired task-specific performance and general task competence limits the broader applicability of these models. This limitation necessitates a shift towards unified models that can concurrently excel at specialized recommendation tasks and maintain their foundational language and reasoning abilities. Mitigating catastrophic forgetting is, therefore, a central research priority. Due to the proprietary nature of LLM training data and the large cost of re-training a conversational agent, we focus on lightweight adaptation methods that can be applied atop a previously trained LLM ∗ Equal contribution. Work done while NV was at Google.
Preprint.
Figure 1: An overview of our O RBIT method. During the fine-tuning of the downstream task, inter-model distance is tracked; when this distance exceeds a threshold ϵ, weight averaging is used as a regularization step to reduce the forgetting of parametric knowledge from θinit . while mitigating forgetting. For this reason, we turn to model merging based methods, which are characterized by 1) their prior success in combining the capabilities of several models directly in parameter space (Yang et al., 2024) and 2) their relatively lightweight nature requiring no substantial retraining. Prior work has used model merging in continual learning settings, where merging serves as an adaptable method to combine prior and current models without inducing substantial forgetting (Marouf et al., 2024; Dziadzio et al., 2025; Kleiman et al., 2025). In investigating the severity of the forgetting problem in GenRetrieval, we find that general LLM text-based reasoning performance is lost early and rapidly during fine-tuning. Relatedly, we find that post-hoc model merging methods designed to boost original model performance after fine-tuning fail to generalize in our setting (Wortsman et al., 2022b). As a result, we focus on model merging techniques that apply merging steps throughout fine-tuning. Motivated by the insufficiencies of one-round post-hoc model merging and by severe forgetting observed in our GenRetrieval setting, we propose O RBIT: Origin-Regulated Back-merging of Intermediate Trajectories. To prevent severe forgetting, O RBIT regulates the total allowable distance between original and fine-tuned model parameters during training, according to specific metrics. At any point when fine-tuned weights are deemed too far, original model parameters are averaged with current fine-tuned parameters. Compared to other regularization techniques, including methods that also employ merging during fine-tuning, we find that O RBIT is Pareto-dominant as measured by recommendation performance and performance on several language-based benchmarks. The key contributions of this work are summarized below: 1. We propose O RBIT, a method that mitigates catastrophic forgetting by tracking inter-model parameter distance and applying weight averaging to constrain drift from the original model. 2. We demonstrate that O RBIT outperforms existing regularization techniques on the GenRetrieval task across multiple task datasets and across multiple text benchmarks. 3. We analyze O RBIT and show that it adopts a distinct averaging schedule from fixed-length repeated merging techniques, which reflects its flexibility and adaptability to different learning behaviors.
2
Related Work
Model merging refers to a set of techniques that combine the capabilities of two or more models by combining their parameters directly in weight space. Wortsman et al. (2022b) introduce a simple method to reduce the forgetting in a pre-trained model after fine-tuning by simply post-hoc interpolating the pre-trained and fine-tuned models. A similar method, LiNeS, is also applied once after fine-tuning, and recombines the task vector resulting from fine-tuning with the pre-trained parameters after rescaling the task vector in a layer-wise manner (Wang et al., 2025). Several techniques introduced in prior work repeatedly apply model merging throughout the training or fine-tuning process. Sanyal et al. (2023) focus only on pretraining a single model, but show that training LLMs from scratch with a high learning rate and intermittent checkpoint averaging can improve generalization. Alexandrov et al. (2024) mitigate forgetting in a multilingual setting by 2
Table 1: Text datasets used for evaluating language and reasoning capabilities. GSM8K MMLU-Pro scoring Acc. 5-shot Yes
Drop
TriviaQA HellaSwag BoolQ ARC-C
sampling Token-F1 1-shot No
sampling Acc. 5-shot No
scoring Acc. 10-shot No
scoring scoring Acc. Acc. 0-shot 25-shot No No
4
Before fine-tuning 0.7 After fine-tuning 0.6 0.5 0.4 0.3 0.2 0.1 0.0 ge BBH SM8K U-Pro Drop iaQA Swag oolQ RC-C era B A G MML Triv Hella t Av x e T
0.35
3
Scores
Recall@10
0.30
2
Recall Text Performance 0.25
1
0.20
0
Tasks
(a) Text benchmark evaluations before and after finetuning Gemma3-1B on the Amazon Review Sports and Outdoors dataset using GenRetrieval. After finetuning, performance drops significantly.
Average Text Performance
BBH
Eval sampling sampling Metric Acc. Acc. N-shot few-shot 8-shot CoT Yes Yes
0
0
200
0
400
Step
0
600
0
800
00
0.15
100
(b) Recall@10 on the Amazon Review Toys and Games dataset and average text performance across our 8 benchmarks during the first 10k steps of baseline GenRetrieval fine-tuning.
Figure 2: Quantitative analysis measuring forgetting during GenRetrieval finetuning. training branched models on different languages interleaved with merging steps. Other work focuses on model merging as a tool to enable continual learning; many prior techniques propose to use model merging to combine models after each task fine-tuning stage in a domain incremental learning setting (Marczak et al., 2024; Marouf et al., 2024; Cheng et al., 2025; Dziadzio et al., 2025). Sokar et al. (2025) extend this type of approach to continual learning across fine-tuning tasks learned with LoRA adapters (Hu et al., 2022). While these prior methods generally apply merging methods after a model has been completely fine-tuned on a new domain, Kleiman et al. (2025) take a slightly different approach by performing model merging sooner after a fixed number of fine-tuning steps. They also consider mitigating forgetting in a single-task adaptation setting, and show their fixed-length averaging scheme also generalizes to this setting. In our work, we consider a single-task adaptation setting, but go beyond fixed-length averaging schedules by merging according to inter-model distance.
3
Background and Motivation
GenRetrieval task In a GenRetrieval recommendation system, the sequential recommendation task is converted to an autoregressive generation problem via framing a user’s item history as a context, and predicting the next item as its completion. Prior work has proposed numerous ways of converting items into token-based IDs, including unstructured and naively structured IDs (Tay et al., 2022), semantically-motivated IDs (Rajput et al., 2023), and learned IDs (Sun et al., 2023); in this work, we focus on the Semantic ID approach for encoding items as proposed in Rajput et al. (2023). In brief, this framework uses a Sentence-T5 model to encode item features (Ni et al., 2022), and then uses an RQ-VAE model to quantize the embedding of the item (Zeghidour et al., 2021). Whereas Rajput et al. (2023) uses an encoder-decoder based model to then learn the GenRetrieval task, but in this work, we adapt pre-trained LLMs for this task. In order to achieve this, we append the Semantic ID token vocabulary to the input and output vocabulary projections of an LLM, and include these new parameters during GenRetrieval fine-tuning. Quantifying the forgetting problem in GenRetrieval While fine-tuning LLMs for GenRetrieval, we noticed the forgetting of original LLM capabilities. To demonstrate this forgetting problem, we compare an instruction-tuned Gemma3-1B model to its GenRetrieval fine-tuned counterpart on text benchmarks (Gemma et al., 2025). For these exemplar experiments, we use the Amazon Product Reviews dataset (He & McAuley, 2016), and report recall on the next item prediction task. 3
Benchmarks to measure original LLM capability can be found in Table 1. As seen in Figure 2a, all text benchmarks drop to levels similar to random chance or majority class performance after GenRetrieval fine-tuning; for example, sampling-based benchmarks BBH, Drop, and TriviaQA drop to 0, and scoring benchmarks ARC-C and BoolQ drop to their majority class performance. Next, we dissect this behavior more closely in our baseline GenRetrieval models to determine when forgetting occurs during fine-tuning. Figure 2b showcases the speed of forgetting during fine-tuning. Within the first 2000 steps of fine-tuning, essentially all text performance is lost, as a ∼0.15 benchmark average reflects a full loss of text performance. This suggests that we need to tailor our mitigation approach towards methods that are well suited for severe and rapid forgetting scenarios.
4
Average Text Performance
Recall@5
One-round merging fails to generalize To Recall 0.030 attempt the re-introduction of general capabil0.35 Text Performance ities back into our GenRetrieval models, we 0.025 experiment with post-hoc weight interpolation 0.30 0.020 between GenRetrieval weights and pretrained LLM weights to improve robustness (Wortsman 0.015 0.25 et al., 2022b; Frankle et al., 2020). We interpo0.010 late the final GenRetrieval and pretrained model 0.20 0.005 parameters with varying interpolation ratios λ (λ = 0 reflects the initial LLM, and λ = 1 re0.000 0.15 flects the GenRetrieval model), and report our 0.0 0.2 0.4 0.6 0.8 1.0 results on text and recall in Figure 3. As seen Interpolation weight by the lack of simultaneous text and recall performance across interpolation ratios, simply avFigure 3: Average text accuracy and Recall@5 eraging the models post fine-tuning is unable performance across post-hoc, one-round weight to recover both sufficient text and retrieval perinterpolations. formance, across interpolation weights. This failure to generalize is likely due to the rapidity of forgetting, as previously observed in Figure 2b. At the time of averaging, enough forgetting has likely already occurred where eventual weight averaging becomes futile. As a result, we turn to techniques that intervene well before fine-tuning is complete.
O RBIT
In this section, we first define preliminaries necessary to define our method, followed by a definition of our method and observations that motivate its design. 4.1
Preliminaries
Notation θinit ∈ Rp refers to initial LLM parameters before GenRetrieval fine-tuning. θcurrent ∈ Rp refers to parameters at the current fine-tuning step. d is a distance function between two sets of parameters, ϵ is a scalar denoting a maximum distance, and T is the total number of training steps. Model distance We use two inter-model distance measures in our work. The first is L2-distance, which is defined as ||θinit − θcurrent ||2 . The second is Sign Dissimilarity (SD), defined as: SD =
# of corresponding parameters with differing signs # of total model parameters
(1)
This distance is inspired in part from a simplification of metrics that measure the “mergeability” of models in prior model merging work (Yadav et al., 2023; Sung et al., 2023). This metric captures the fraction of parameters that have undergone a meaningful change, as measured by sign flipping, a property that is efficient to compute via bitwise XOR. This also means that only sign bits need to be stored for computation. Since SID vocabulary parameters are randomly initialized for fine-tuning in θinit , we exclude them from distance computations. 4
4.2
Algorithm
Given the early and rapid degradation of text performance during fine-tuning, we are interested in a method that intervenes quickly during this process. While other methods that also employ merging during fine-tuning may be able to intervene quickly, these methods generally use a fixed cadence throughout training, which may not be the optimal schedule for the entire training duration (Kleiman et al., 2025). Given these observations, we propose O RBIT: Origin-Regulated Back-Merging of Intermediate Trajectories. Our method is simple yet surprisingly effective: for a given metric, we fix a maximum distance allowable between the original model parameters, which serves as our origin, and the model parameters. When a training step causes the current model parameters to exceed this distance, we average the original parameters with the offending current model parameters (just after the gradient update), which we refer to as “back-merging”. This method schedules averaging as a function of inter-model distance, providing two major benefits. The first is its distance guarantee, where the trained model is regulated to be within a fixed distance of the original model. The second is the added flexibility of the averaging schedule. This flexibility allows for distance to dictate the averaging schedule, allowing the method to trigger merges as needed, which in turn can help improve generalizability. In brief, O RBIT regularizes fine-tuning via tracking a distance between θcurrent and θinit , and triggering a weight averaging step if this distance is deemed too large. We summarize O RBIT in Algorithm 1. Algorithm 1 O RBIT: Origin-Regulated Back-merging of iterative trajectories 1: Input: θinit , T, ϵ, d(·, ·) 2: for t in 1, ..., T do ∗ 3: θt+1 = θt − η∇θt Ltask {Optimizer Update} ∗ 4: while d(θt+1 , θinit ) > ϵ then θ∗
+θinit
∗ 5: θt+1 = t+12 ∗ 6: θt+1 = θt+1 7: end for 8: return θT
{Back-merging}
A potential concern with using Sign Dissimilarity (SD) as the trigger metric in O RBIT is that unlike L2 distance, SD does not contract uniformly with an arithmetic averaging step. A sign flip at coordinate i survives averaging when |θcurrent (i)| > |θinit (i)|. We show that this does not affect our within-distance guarantee in O RBIT.2 Crucially, if averaging fails to bring SD below the threshold ϵ, the next post-update check re-triggers another merge. We demonstrate that the resulting sequence of consecutive merges drives SD to zero in a bounded number of steps. This suffices to show that SD can be driven below some threshold ϵ within a bounded number of steps. Our proof is in Appendix D. 4.3
Why use inter-model distance?
While averaging as a regularization tool is inspired by prior work (Kleiman et al., 2025; Marczak et al., 2024; Marouf et al., 2024), the use of inter-model distance to schedule averaging steps is an important and novel distinction in our work. Our motivation to use distance to dictate averaging steps comes from 1) a longstanding notion of model knowledge corresponding to localities in weight space (Gueta et al., 2023; Wortsman et al., 2022a) and 2) a preliminary study of different checkpoints produced by Soup-to-Go fine-tuning. Soup-to-Go is a regularization method where for a fixed number of steps T , the authors define a hyperparameter 0 < p < 1 where averaging occurs every pT steps between θinit and θcurrent (Kleiman et al., 2025). In this analysis, we choose Sign Dissimilarity as our distance measure, and we compute the text performance of these checkpoints. We display the results of this analysis in Figure 4, where we evaluate several checkpoints (every 2000 steps) from fine-tuning Gemma1B IT on the Amazon Product Reviews Sports and Outdoors dataset using GenRetrieval. We use Soup-to-Go to regularize this fine-tuning by averaging every 1000 steps. Additional hyperparameters are in Table 7 in Appendix C. In testing the performance of Soup-to-Go in our setting, we observe that across checkpoints saved throughout training, there is a correlation between text performance and distance from the initial 2 Despite allowing repeated merges, we note empirically this is never needed in our experiments.
5
Text Performance
0.40
Fit (R²=0.83)
0.35 0.30 0.25 0.20 0.002 0.004 0.006 0.008 0.010 0.012
Sign Dissimilarity (SD)
Figure 4: A scatter plot demonstrating the correlation between sign dissimilarity (SD) and average text performance. Points are collected from a Soup-to-Go experiment with a cadence of 1000 steps.
Figure 5: LLM-based GenRetrieval setup with separate text and SID vocabularies. Such a model can handle GenRetrieval queries with both SID tokens and text control tokens, and general LLM queries composed of solely text tokens.
parameters. This observation suggests that fine-tuning that pushes the model starkly away from its starting point can induce severe forgetting. As a result, we focus on limiting this distance in our proposed method in order to preserve capability from the original model. In this setting, inter-model distance serves as a lightweight, gradient-free, and data-free proxy to forgetting.
5
Experimental Setup
5.1
GenRetrieval Fine-tuning
For fine-tuning our GenRetrieval models, we use with Gemma3 (Gemma et al., 2025) as our base model.3 We use the instruction-tuned version as it more closely matches the capabilities we are interested in preserving versus the pre-trained model. For our recommendation data, we use the Amazon Product Reviews dataset, which comprises three different subsets: Beauty, Sports and Outdoors, and Toys and Games (He & McAuley, 2016). Dataset statistics are in Appendix B. We follow Rajput et al. (2023) for dataset preprocess- Table 2: Base GenRetrieval Hyperparameters ing steps, including filtering users with less than 5 reviews, and limiting the number of items in a user’s Hyperparameter Value history to 20. After preprocessing, each datapoint Optimizer Adafactor consists of a user ID, the user’s previous items enLR Schedule Cosine Decay coded with Semantic IDs (SIDs), and a held-out addiPeak Learning Rate 0.02 tional item which is also converted to its SIDs. Each Min. Learning Rate 1e-5 item uses 4 SID tokens to create the SID sequence. Warmup Steps 10,000 Decay Steps 30,000 We add <start_of_SID> and <end_of_SID> Training Steps 50,000 control tokens to the set of tokens in our base LLM, Batch Size 16 and append these tokens to each prompt in order to elicit the SID response during evaluation. To train the baseline GenRetrieval models, we logically separate SID vocabulary items from the original text vocabulary in order to compute cross-entropy loss from both text and SID vocabularies. SID parameters are also excluded from back-merging. An example of our model and an example GenRetrieval query can be found in Figure 5. Note that for GenRetrieval finetuning, the text cross-entropy is not used beyond control tokens since the decoded tokens only contain the SID tokens. During inference, we default to text vocab and only use SID vocab for tokens between <start_of_SID> and <end_of_SID> control tokens.
To compute the possible SID sequences for computing retrieval metrics, we use beam search with 20 beams, and 20 tokens per beam. To evaluate sequential recommendation performance, we report 3 Gemma3 is released under the Gemma Terms of Use.
6
Normalized Discounted Cumulative Gain@10 (NDCG@10) and Recall@10 values. We summarize key hyperparameters for training our baseline GenRetrieval models in Table 2. We use a cosine decay learning rate with warmup for these baselines. 5.2
Baselines
Simple baselines We include no-intervention baselines, as well as L2-weight decay to reflect a baseline from traditional continual learning literature. Other techniques in this space, like Elastic Weight Consolidation and data replay methods, require data from the original model training data to reintroduce into the training mixture or compute Fisher information (Kirkpatrick et al., 2017). In our setting, we assume no access to prior training data, as we are interested in preserving the capability of LLMs trained on proprietary data. (1) No interventions refer to baseline models without any training interventions. The first of these models is simply Gemma3-1B-IT, which represents full text capability, and no GenRetrieval performance. The second baseline is the fine-tuned GenRetrieval model. (2) L2 weight decay is a classic continual learning technique that adds a loss penalty to minimize the distance between the current and initial model weights. Our proposed O RBIT method is similar in spirit to weight decay, but sets a strict boundary for model distance and averages parameters accordingly rather than learning according to a penalty. Soup-to-Go Soup-to-Go is a simple method designed for continual learning in deep-learning models (Kleiman et al., 2025). While fine-tuning a model on a new domain dataset the original model weights are averaged with the current model weights after every pT steps in order to preserve capabilities from the original model, where 0 < p < 1, and T is the total number of training steps. The authors generally use < 10 merges during fine-tuning for their experiments, with p > 0.1. In this work, we specify Soup-to-Go baselines by their cadence with k steps. 5.3
O RBIT Settings
For both Soup-to-Go and O RBIT, we utilize a constant learning rate during fine-tuning. While a cosine decay schedule helps the GenRetrieval baseline method improve recall performance, a constant learning rate simplifies the effect of model averaging in Soup-to-Go and O RBIT, so different learning segments between merges are not subject to different learning rates. In turn, we fine-tune these models longer, for up to 200k steps. Longer training was not found to be helpful in our baseline GenRetrieval models. We use a learning rate of 0.001 for all temporal averaging based experiments. We test ϵ = 7e − 3, 7.5e − 3 given the results in Figure 4. 5.4
Text evaluation
To evaluate the language capabilities of our GenRetrieval models, we use the benchmarks from Table 1, as described with their evaluation settings. These benchmarks reflect a mix of scoring and sampling-based evaluations in order to cover a breadth of capabilities. We report the average.
6
Results
We evaluate our baselines and O RBIT for GenRetrieval fine-tuning across two primary objectives: textbased reasoning and sequential recommendation. We report average text performance, NDCG@10, and Recall@10. To identify techniques that perform well in both domains, we employ Pareto efficiency principles. Because Pareto-optimal sets often contain multiple solutions, we introduce a Distance To Ideal Point (DTIP) metric to select a single, balanced model. DTIP measures the normalized Euclidean distance between a model’s performance as represented by a tuple, and an "ideal" point representing the theoretical maximum in both domains. For a given model, let T represent text performance, and R represent retrieval performance (Recall@5). We first normalize the scores using the min-max scaling to a range [0, 1]. T − Tmin R − Rmin T′ = and R′ = (2) Tmax − Tmin Rmax − Rmin Minimum text performance is set by the text performance observed during full, no-intervention GenRetrieval fine-tuning. Maximum performance is set by the specialized, single-domain models. 7
Table 3: Full results across Amazon Product Reviews test sets for baseline methods and O RBIT. To identify a single representative checkpoint from a Pareto-optimal set, we measure the fraction of original performance retained for both text and retrieval. We then select the checkpoint that maximizes this combined retention. Bold is best performance, underline is second best. Sports and Outdoors Method
Avg Text Perf.
NDCG @10
Recall @10
Text Baseline Retrieval Baseline
35.72 15.52
0 2.16
L2 Decayλ=1e−4 Soup-to-Gok=3K Soup-to-Gok=2K
15.76 26.73 30.27
O RBITSD=7.5e−3 O RBITSD=7e−3
28.95 28.88
Toys and Games DTIP↓
Avg Text Perf.
NDCG @10
Recall @10
0 3.76
1.00 1.00
35.72 15.59
0 3.98
1.92 1.20 1.11
3.53 2.32 2.15
0.99 0.63 0.54
15.66 20.34 28.55
1.37 1.22
2.58 2.42
0.49 0.55
30.86 30.98
Beauty Subsets DTIP↓
Avg Text Perf.
NDCG @10
Recall @10
DTIP↓
0 6.57
1.00 1.00
35.72 15.69
0 3.47
0 6.06
1.00 1.00
3.67 2.42 2.54
6.09 4.26 4.36
1.00 0.86 0.50
15.75 32.70 26.44
3.17 0.79 1.88
5.49 1.63 3.64
1.00 0.78 0.63
2.64 2.59
4.55 4.40
0.41 0.40
25.33 30.47
2.28 2.11
4.37 3.90
0.62 0.47
The DTIP is then calculated as the L2 distance from the normalized performance (T ′ , R′ ) to the ideal point, which is (1, 1), reflecting maximum text and retrieval performance. DT IP = ||(1, 1) − (T ′ , R′ )||2 6.1
(3)
O RBIT maximizes joint performance
We measure the performance across the LLM benchmarks from Section 5.4 and on the sequential recommendation task for each chosen subset of the Amazon Product Reviews dataset. For both O RBIT and Soup-to-Go fine-tuning, we save a set of checkpoints during training that are Pareto-optimal. To select a single representative checkpoint from this Pareto-optimal set, we measure the DTIP, and then select the checkpoint that maximizes this combined retention. We compare the performance of O RBIT to several baselines and report our results in Table 3. We find that O RBIT improves joint performance further than Soup-to-Go, minimizing the distance to the ideal point. Beyond comparisons to baselines, we note that performance of a ∼0.3 text average on text tasks and Recall@5 of ∼.02 reflects substantial text and retrieval performance, compared to their topline baselines. We also display the set of Pareto-optimal points for these methods in Figure 6, specifically for the Sports and Outdoors validation subset. As seen in the figure, O RBIT improves joint performance even further than Soup-to-Go, with all four displayed variations outperforming all baselines. We note that in comparing all Pareto optimal checkpoints per training run, we see that every checkpoint generated by O RBIT is superior to every checkpoint generated by baselines, which expands upon the results found in the table. We hypothesize that the imposition of model distance to determine averaging helps our method find solutions that maintain high text performance by design, while optimizing the GenRetrieval objective subject to this soft constraint. 6.2
Temporal averaging techniques mitigate text forgetting in GenRetrieval
As seen in both Figure 6 and Table 3, only the methods that use repeated averaging, namely Soup-toGo and O RBIT, are able to achieve non-trivial text and retrieval performance. This result emphasizes the utility of regularization techniques that employ weight averaging multiple times during training; in the face of severe forgetting, classical techniques like weight decay, as well as more modern techniques like post-hoc weight averaging, may fail to generalize. For Soup-to-Go results, we select k = 2000, 3000 as they provides an adequate balance between text and retrieval performance. However, we note that this value, which corresponds to p = 0.01, 0.015 according to the original definition of the method, is much smaller than the values reported in the original paper, which are all larger than 0.15. This finding highlights the severe degree of forgetting, and the greater importance of repeated averaging strategies in fine-tuning LLMs for GenRetrieval.
7
Analysis
7.1
Our O RBIT schedule is distinct from a constant schedule
In designing O RBIT, we hypothesized that the flexibility in the learned schedule, as determined by inter-model distance, may help prevent forgetting compared to a fixed-length averaging schedule. To 8
Text Task Average
0.35 0.30 Text Baseline Retrieval Baseline L2 Decay Soup-to-Go, k=2000 ORBIT SD=6.5e-3 ORBIT SD=7e-3 ORBIT SD=7.5e-3 ORBIT SD=8e-3
0.25 0.20 0.15
0.000 0.005 0.010 0.015 0.020 0.025 0.030
Recall@5
Figure 6: Text and recall performance for O RBIT models, compared to a Soup-to-Go baseline and L2 decay baselines on the Sports and Outdoors dataset (validation) and our 8 text benchmarks. We display only Pareto-optimal checkpoints generated within each experiment. We can observe that all O RBIT checkpoints outperform those generated from Soup-to-Go training.
Figure 7: The number of steps between averaging events (indexed) in O RBIT with a maximum sign dissimilarity of 0.007, over 200k training steps. The number of steps between averages increases during training, before settling around 3000 steps later in training.
observe the learned schedule, we compute the number of training steps between each averaging step for O RBIT applied to GenRetrieval with the Sports and Outdoors dataset. As seen in Figure 7, the number of steps between averaging steps increases over the training interval, before appearing to converge to about 3000 steps in the latter half of training. This “learned” schedule is distinct from what is proposed in the original Soup-to-Go work, where averaging occurs at a regular interval, which would be reflected as a horizontal line in this graph. This added flexibility inherent in our method may help its generalization as well to other settings. 7.2
Choice of Metric
Table 4: Metric analysis of SD and L2 in O RBIT. We find that SD is preferred to L2 distance for O RBIT.
Table 5: O RBIT and Soup-to-Go applied to GenRetrieval fine-tuning, using Gemma3-4B IT as the base model.
Setting
Avg Text Perf.
NDCG @10
Recall @10
DTIP↓
Setting
Avg Text Perf.
NDCG @10
Recall @10
DTIP ↓
O RBITSD=0.007 O RBITL2=5 O RBITL2=50 O RBITL2=500
28.88 35.40 35.66 15.82
1.22 0.25 0.71 1.36
2.42 0.36 1.26 2.58
0.55 0.88 0.67 1.05
Gemma3-4B IT Base GenRetrieval Soup-to-Gok=2K O RBITSD=0.007
57.02 15.56 48.12 47.94
0 2.04 1.32 1.35
0 3.63 2.54 2.58
1.00 1.00 0.43 0.42
We initially select Sign Dissimilarity as the metric in O RBIT as it provides some indication of how many parameters change meaningfully, versus a combined magnitude of change, as is the case with L2 distance. However, we are interested in the sensitivity of O RBIT to different choices of distance metrics. To this end, we compare the use of SD and L2 in O RBIT; for SD, we fix the distance to be 0.007 given its success in prior experiments, and for L2, we test values of 5, 50, 500. We display our result on Sports and Outdoors in Table 4. We find that SD is preferred to L2 distance, but that L2 distance may also serve as a suitable metric. However, we recommend the use of SD given its simplicity to compute, as well as its fractional representation of parameter change. 7.3
Scaling
While we focus on experimenting with Gemma3-1B models, we are interested in evaluating the performance of our method on a larger model to determine the method’s extensibility. We evaluate both Soup-to-Go and O RBIT on Gemma-3-4B models, on the Sports and Outdoors dataset. We select 9
hyperparameters for both methods given their best performance at the 1B scale. We display our results on Gemma3-4B IT in Table 5. We find that performance scales with a larger model, which provides improved overall performance.
8
Conclusion
In this work, we introduce O RBIT, a new method that enables language models to perform both general text and GenRetrieval functionalities. We design O RBIT to prevent LLMs from losing their general language skills while being fine-tuned for the specialized GenRetrieval task. We demonstrate that O RBIT preserves a meaningful amount of both recommendation and text-based reasoning capabilities in LLMs adapted for the GenRetrieval task, which in turn can help enable a unified model for conversational based discovery for recommendation systems. We also demonstrate that compared to both classic regularization methods and related averaging-based methods, O RBIT outperforms these methods by maintaining higher Pareto performance across both desired capabilities. While our work focuses on mitigating forgetting specific to GenRetrieval, our method is general in its construction and could extend to other applications where severe forgetting is observed.
References Alexandrov, A., Raychev, V., Müller, M. N., Zhang, C., Vechev, M., and Toutanova, K. Mitigating catastrophic forgetting in language transfer via model merging. arXiv preprint arXiv:2407.08699, 2024. Cheng, F., Wang, Z., Sung, Y.-L., Lin, Y.-B., Bansal, M., and Bertasius, G. Dam: Dynamic adapter merging for continual video qa learning. In 2025 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), pp. 6805–6817. IEEE, 2025. Dziadzio, S., Udandarao, V., Roth, K., Prabhu, A., Akata, Z., Albanie, S., and Bethge, M. How to merge your multimodal models over time? In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 20479–20491, 2025. Frankle, J., Dziugaite, G. K., Roy, D., and Carbin, M. Linear mode connectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pp. 3259–3269. PMLR, 2020. Gemma, T., Kamath, A., Ferret, J., Pathak, S., Vieillard, N., Merhej, R., Perrin, S., Matejovicova, T., Ramé, A., Rivière, M., et al. Gemma 3 technical report. arXiv preprint arXiv:2503.19786, 2025. Gueta, A., Venezian, E., Raffel, C., Slonim, N., Katz, Y., and Choshen, L. Knowledge is a region in weight space for fine-tuned language models. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 1350–1370, 2023. He, R. and McAuley, J. Ups and downs: Modeling the visual evolution of fashion trends with one-class collaborative filtering. In proceedings of the 25th international conference on world wide web, pp. 507–517, 2016. Hu, E. J., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., Chen, W., et al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. Kirkpatrick, J., Pascanu, R., Rabinowitz, N., Veness, J., Desjardins, G., Rusu, A. A., Milan, K., Quan, J., Ramalho, T., Grabska-Barwinska, A., et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114(13):3521–3526, 2017. Kleiman, A., Dziugaite, G. K., Frankle, J., Kakade, S., and Paul, M. Soup to go: mitigating forgetting during continual learning with model averaging. arXiv preprint arXiv:2501.05559, 2025. Marczak, D., Twardowski, B., Trzciński, T., and Cygert, S. Magmax: Leveraging model merging for seamless continual learning. In European Conference on Computer Vision, pp. 379–395. Springer, 2024. Marouf, I. E., Roy, S., Tartaglione, E., and Lathuilière, S. Weighted ensemble models are strong continual learners. In European Conference on Computer Vision, pp. 306–324. Springer, 2024. 10
Ni, J., Abrego, G. H., Constant, N., Ma, J., Hall, K., Cer, D., and Yang, Y. Sentence-t5: Scalable sentence encoders from pre-trained text-to-text models. In Findings of the Association for Computational Linguistics: ACL 2022, pp. 1864–1874, 2022. Rajput, S., Mehta, N., Singh, A., Hulikal Keshavan, R., Vu, T., Heldt, L., Hong, L., Tay, Y., Tran, V., Samost, J., Kula, M., Chi, E., and Sathiamoorthy, M. Recommender systems with generative retrieval. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 10299–10315. Curran Associates, Inc., 2023. URL https://proceedings.neurips.cc/paper_files/paper/2023/ file/20dcab0f14046a5c6b02b61da9f13229-Paper-Conference.pdf. Sanyal, S., Neerkaje, A., Kaddour, J., Kumar, A., and Sanghavi, S. Early weight averaging meets high learning rates for llm pre-training. arXiv preprint arXiv:2306.03241, 2023. Sokar, G., Dziugaite, G. K., Arnab, A., Iscen, A., Castro, P. S., and Schmid, C. Continual learning in vision-language models via aligned model merging. arXiv preprint arXiv:2506.03189, 2025. Sun, W., Yan, L., Chen, Z., Wang, S., Zhu, H., Ren, P., Chen, Z., Yin, D., Rijke, M., and Ren, Z. Learning to tokenize for generative retrieval. Advances in Neural Information Processing Systems, 36:46345–46361, 2023. Sung, Y.-L., Li, L., Lin, K., Gan, Z., Bansal, M., and Wang, L. An empirical study of multimodal model merging. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 1563–1575, 2023. Tay, Y., Tran, V., Dehghani, M., Ni, J., Bahri, D., Mehta, H., Qin, Z., Hui, K., Zhao, Z., Gupta, J., et al. Transformer memory as a differentiable search index. Advances in Neural Information Processing Systems, 35:21831–21843, 2022. Wang, K., Dimitriadis, N., Favero, A., Ortiz-Jimenez, G., Fleuret, F., and Frossard, P. Lines: Post-training layer scaling prevents forgetting and enhances model merging. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/ forum?id=J5sUOvlLbQ. Wortsman, M., Ilharco, G., Gadre, S. Y., Roelofs, R., Gontijo-Lopes, R., Morcos, A. S., Namkoong, H., Farhadi, A., Carmon, Y., Kornblith, S., et al. Model soups: averaging weights of multiple finetuned models improves accuracy without increasing inference time. In International conference on machine learning, pp. 23965–23998. PMLR, 2022a. Wortsman, M., Ilharco, G., Kim, J. W., Li, M., Kornblith, S., Roelofs, R., Lopes, R. G., Hajishirzi, H., Farhadi, A., Namkoong, H., et al. Robust fine-tuning of zero-shot models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 7959–7971, 2022b. Yadav, P., Tam, D., Choshen, L., Raffel, C. A., and Bansal, M. Ties-merging: Resolving interference when merging models. Advances in Neural Information Processing Systems, 36:7093–7115, 2023. Yang, E., Shen, L., Guo, G., Wang, X., Cao, X., Zhang, J., and Tao, D. Model merging in llms, mllms, and beyond: Methods, theories, applications and opportunities. arXiv preprint arXiv:2408.07666, 2024. Zeghidour, N., Luebs, A., Omran, A., Skoglund, J., and Tagliasacchi, M. Soundstream: An end-to-end neural audio codec. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 30: 495–507, 2021.
11
A
Limitations
This work focuses on mitigating forgetting in a full fine-tuning setting, common in model post-training pipelines. It does not consider parameter efficient fine-tuning (PEFT) techniques such as LoRA, which may have different forgetting behaviors during fine-tuning (Hu et al., 2022). Additionally, we test our method on 1B and 4B models, which are much smaller than frontier models. Testing whether O RBIT achieves the same forgetting reduction at this scale remains necessary.
B
Amazon Product Reviews Dataset
Table 6: Dataset Statistics for the three subsections of the Amazon Product Reviews data. Dataset Users Items Beauty Sports and Outdoors Toys and Games
C
22,363 35,598 19,412
12,101 18,357 11,924
Soup-to-Go exploratory hyperparameters Table 7: Soup-to-Go Hyperparameters in Distance Study Hyperparameter
Value
Optimizer LR Schedule Peak Learning Rate Min. Learning Rate Warmup Steps Decay Steps Training Steps Batch Size
Adafactor Cosine Decay 0.02 1e-5 10,000 20,000 30,000 16
D
Finite-Merge Recovery Guarantee for ORBIT under Sign Dissimilarity
D.1
Setup
Let θinit ∈ Rp denote the origin parameters and let θ0 ∈ Rp denote the current parameters during fine-tuning where O RBIT triggers a merge. Define the merge operator M (θ) = 21 θ + θinit ,
(4)
and let θk = M k (θ0 ) denote the parameters obtained by applying k consecutive merges. We make one mild non-degeneracy assumption: θinit (i) ̸= 0 for every coordinate i counted in SD. The randomly-initialized SID vocabulary entries are explicitly excluded from the SD computation, as discussed in Section 4.1, and pretrained weights are generically nonzero. D.2
Closed form for iterated merging
For all k ≥ 0, θk =
1 1 θ0 + 1 − k θinit . k 2 2 12
(5)
Proof. By induction on k. The base case k = 0 is immediate. Assuming the formula holds at step k, 1 θk + θinit 2 1 1 1 θ + (1 − )θ + θ = 0 init init 2 2k 2k 1 1 = k+1 θ0 + 1 − k+1 θinit . 2 2
θk+1 =
D.3
(6) (7)
Per-coordinate sign flip
For a coordinate i with θinit (i) ̸= 0, and suppose sign(θ0 (i)) ̸= sign(θinit (i)). Define the magnitude ratio |θ0 (i)| ri = . (8) |θinit (i)| Then sign flipping after k merges will occur, sign(θk (i)) = sign(θinit (i)), if and only if 2k > 1 + ri .
(9)
Proof. By Result D.2, θk (i) = 21k θ0 (i) + (1 − 21k )θinit (i). Without loss of generality assume θinit (i) > 0. By the assumption, θ0 (i) < 0. Then θk (i) > 0 iff 1 1 θ0 (i) + (1 − k )θinit (i) > 0 2k 2 (2k − 1) θinit (i) > −θ0 (i)
(10) (11)
k
⇐⇒ 2 > 1 + ri . Coordinates that are not flipped at θ0 remain aligned with θinit for every k ≥ 0, since averaging two same-sign values cannot change the sign. D.4
Finite-merge recovery
Let F = { i : sign(θ0 (i)) ̸= sign(θinit (i)) } denote the set of initially-flipped coordinates, and let rmax = maxi∈F ri . For any k > log2 (1 + rmax ), SD(θk , θinit ) = 0.
(12)
Proof. By Result D.3, the set of coordinates contributing to SD(θk , θinit ) at merge step k is exactly Fk = i ∈ F : ri ≥ 2k − 1 , (13) which is monotonically nonincreasing in k. Choose k so that 2k − 1 > rmax , i.e., k > log2 (1 + rmax ). Then Fk = ∅ and SD(θk , θinit ) = 0.
13