Your Language Model May Think Too Rigidly: Achieving Reasoning Consistency with Symmetry-Enhanced Training

TL;DR

This paper introduces a data augmentation method called MEND that enhances large language models' reasoning consistency by improving their robustness to query variations through symmetry-aware training, leading to better generalization.

Contribution

The paper proposes a novel symmetry-enhanced data augmentation technique that improves LLM robustness and reasoning performance across varied query phrasings, focusing on knowledge extraction stages.

Findings

MEND improves reasoning accuracy across diverse query variations.

The approach enhances model robustness to out-of-distribution data.

Experiments show better generalization in logical and arithmetic reasoning tasks.

Abstract

Large Language Models (LLMs) have demonstrated strong reasoning capabilities across various tasks. However, even minor variations in query phrasing, despite preserving the underlying semantic meaning, can significantly affect their performance. To address this, we focus on enhancing LLMs' awareness of symmetry in query variations and propose syMmetry-ENhanceD (MEND) Data Augmentation, a data-centric approach that improves the model's ability to extract useful information from context. Unlike existing methods that emphasize reasoning chain augmentation, our approach improves model robustness at the knowledge extraction stage through query augmentations, enabling more data-efficient training and stronger generalization to Out-of-Distribution (OOD) settings. Extensive experiments on both logical and arithmetic reasoning tasks show that MEND enhances reasoning performance across diverse query variations, providing new insight into improving LLM robustness through structured dataset curation.