Leveraging parameter space symmetries for reasoning skill transfer in LLMs
Utilizing an alignment-first strategy to transfer advanced reasoning skills to a non-reasoning model.
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.
Latest publications
GPO: learning from critical steps to improve LLM reasoning
A novel fine-tuning strategy designed to improve LLM multi-step reasoning capabilities by focusing on pivotal moments.
NeurIPSViCrit: a verifiable reinforcement learning proxy task for visual perception in VLMs
An RL proxy task that trains VLMs to localize synthetic hallucinations injected into human-written captions.
NeurIPSInfluence functions for efficient data selection in reasoning
A proposal to define reasoning data quality using influence functions.
NeurIPS