Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Leveraging Parameter Space Symmetries for Reasoning Skill Transfer in LLMs

Stefan Horoi, Sangwoo Cho, Supriyo Chakraborty, Shi-Xiong Zhang, Sambit Sahu, Guy Wolf, Genta Indra Winata

TL;DR

This work tackles the problem of transferring hard-earned reasoning skills between evolving LLM families while avoiding negative interference. It proposes an alignment-first approach that exploits transformer parameter-space symmetries—permutation for SwiGLU and rotation/scaling for GQA—to align source and target models before applying task arithmetic. By extending these alignment techniques to modern architectures and evaluating on challenging reasoning benchmarks, the method achieves substantial gains over standard task arithmetic and enables skill transfer across divergent development paths. This reduces the need for full re-finetuning and enhances adaptability of specialized skills across model versions, with practical impact for rapid, efficient skill merging in evolving LLM ecosystems.

Abstract

Task arithmetic is a powerful technique for transferring skills between Large Language Models (LLMs), but it often suffers from negative interference when models have diverged during training. We address this limitation by first aligning the models' parameter spaces, leveraging the inherent permutation, rotation, and scaling symmetries of Transformer architectures. We adapt parameter space alignment for modern Grouped-Query Attention (GQA) and SwiGLU layers, exploring both weight-based and activation-based approaches. Using this alignment-first strategy, we successfully transfer advanced reasoning skills to a non-reasoning model. Experiments on challenging reasoning benchmarks show that our method consistently outperforms standard task arithmetic. This work provides an effective approach for merging and transferring specialized skills across evolving LLM families, reducing redundant fine-tuning and enhancing model adaptability.

Leveraging Parameter Space Symmetries for Reasoning Skill Transfer in LLMs

TL;DR

This work tackles the problem of transferring hard-earned reasoning skills between evolving LLM families while avoiding negative interference. It proposes an alignment-first approach that exploits transformer parameter-space symmetries—permutation for SwiGLU and rotation/scaling for GQA—to align source and target models before applying task arithmetic. By extending these alignment techniques to modern architectures and evaluating on challenging reasoning benchmarks, the method achieves substantial gains over standard task arithmetic and enables skill transfer across divergent development paths. This reduces the need for full re-finetuning and enhances adaptability of specialized skills across model versions, with practical impact for rapid, efficient skill merging in evolving LLM ecosystems.

Abstract

Task arithmetic is a powerful technique for transferring skills between Large Language Models (LLMs), but it often suffers from negative interference when models have diverged during training. We address this limitation by first aligning the models' parameter spaces, leveraging the inherent permutation, rotation, and scaling symmetries of Transformer architectures. We adapt parameter space alignment for modern Grouped-Query Attention (GQA) and SwiGLU layers, exploring both weight-based and activation-based approaches. Using this alignment-first strategy, we successfully transfer advanced reasoning skills to a non-reasoning model. Experiments on challenging reasoning benchmarks show that our method consistently outperforms standard task arithmetic. This work provides an effective approach for merging and transferring specialized skills across evolving LLM families, reducing redundant fine-tuning and enhancing model adaptability.

Paper Structure

This paper contains 17 sections, 9 equations, 1 figure, 3 tables.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Preliminaries and Methodology
  4. Task Arithmetic for Skill Transfer
  5. Leveraging Parameter Space Symmetries in Neural Networks
  6. Aligning SwiGLU Layers via Permutation Symmetry
  7. Aligning GQA Layers via Rotation and Scaling Symmetry
  8. Optimal Rotation for Query and Key Heads
  9. Optimal Scaling for Query and Key Heads
  10. Results
  11. Experimental Setup
  12. Parameter Space Alignment Improves Reasoning Skill Transfer
  13. Conclusion
  14. Experimental details
  15. Data used for activation-based alignment
  16. ...and 2 more sections

Figures (1)

  • Figure 1: The family tree of models used in our experiments. The diagram illustrates the two divergent fine-tuning branches originating from the common ancestor, Llama-3.1-Base. The top branch leads to the reasoning-specialized Nemotron-Nano built on top of Llama-3.1-Instruct, while the bottom branch shows the parallel instruction-tuning of the Tulu3 model. Our work aims to transfer reasoning skills acquired on the top branch to the models on the parallel bottom branch by first aligning the parameter space of Tulu3 to the one of Llama-3.1-Instruct, and then adding the reasoning skill vector.