LayAlign: Enhancing Multilingual Reasoning in Large Language Models via Layer-Wise Adaptive Fusion and Alignment Strategy

TL;DR

LayAlign tackles the challenge of multilingual reasoning in English-centric LLMs by layer-wise fusing representations from all layers of a multilingual encoder and integrating them into a decoder-only LLM via a per-layer adaptive fusion and cross-attention mechanism. The method freezes both the multilingual encoder and the LLM, training only a lightweight adapter and a layer-wise aligner in two stages to align representation spaces and improve cross-lingual reasoning. Across mathematical reasoning, commonsense reasoning, and language understanding tasks, LayAlign consistently outperforms strong baselines and state-of-the-art approaches, with notable gains for low-resource languages and robust cross-lingual generalization. The work also provides thorough analyses of representation space, encoder choices, and input strategies, showing that aligning multilingual signals into a shared space enhances multilingual reasoning with practical implications for expanding access to AI capabilities in low-resource languages.

Abstract

Despite being pretrained on multilingual corpora, large language models (LLMs) exhibit suboptimal performance on low-resource languages. Recent approaches have leveraged multilingual encoders alongside LLMs by introducing trainable parameters connecting the two models. However, these methods typically focus on the encoder's output, overlooking valuable information from other layers. We propose \aname (\mname), a framework that integrates representations from all encoder layers, coupled with the \attaname mechanism to enable layer-wise interaction between the LLM and the multilingual encoder. Extensive experiments on multilingual reasoning tasks, along with analyses of learned representations, show that our approach consistently outperforms existing baselines.

Figures (9)

• Figure 1: Overview of LayAlign. A multilingual encoder is aligned with the target LLM with an adapter and the layer-wise aligner. We keep the multilingual encoder and LLM frozen, whereas the adapter and layer-wise aligner are optimized in two stages. For simplicity, shifted output tokens were omitted from the input representation. Left: (1) Translation stage. In this stage, LayAlign is fine-tuned using translation data, where the data consists of translations from other languages into English. Right: (2) Task Stage. In this stage, LayAlign is fine-tuned using specialized downstream task data, where the input is multilingual and the output is in English.
• Figure 2: The illustration of our proposed Adaptive Fusion-Enhanced Attention. It consists of self-attention (right), cross-attention (left), and a gate module. Both cross-attention and self-attention modules share the same linear weights as that of the backbone LLM.
• Figure 3: Experimental results for the Swahili language on the MGSM and MSVAMP datasets.
• Figure 4: The cosine similarities of the final layer of LLM pooled output representations of English with other languages obtained with the FLORES-101 dataset.
• Figure 5: First two principal components of pooled output representations obtained with the FLORES-101.