Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Gated Tree Cross-attention for Checkpoint-Compatible Syntax Injection in Decoder-Only LLMs

Xinyu Gao, Shaonan Wang, Nai Ding

TL;DR

This work tackles the brittleness of decoder-only LLMs to syntactic perturbations by introducing GTCA, a checkpoint-compatible gated tree cross-attention pathway that reads cached constituency chunk memory without altering the backbone. A token update mask and a three-stage training schedule regulate the structural updates, enabling targeted syntax improvements while preserving broad competence across tasks and backbones. Empirical results on CLOTH, MMLU, BLiMP, CoLA, HellaSwag, and Winogrande show consistent syntactic gains with minimal or no loss in MCQA performance, and layer-wise probing confirms syntax-aligned internal changes. The approach remains practical by using offline parsing, memory caching, and a gating mechanism that provides interpretability and stability, though it relies on external parsers and incurs preprocessing/compute overhead.

Abstract

Decoder-only large language models achieve strong broad performance but are brittle to minor grammatical perturbations, undermining reliability for downstream reasoning. However, directly injecting explicit syntactic structure into an existing checkpoint can interfere with its pretrained competence. We introduce a checkpoint-compatible gated tree cross-attention (GTCA) branch that reads precomputed constituency chunk memory while leaving backbone architecture unchanged. Our design uses a token update mask and staged training to control the scope and timing of structural updates. Across benchmarks and Transformer backbones, GTCA strengthens syntactic robustness beyond continued-training baselines without compromising Multiple-Choice QA performance or commonsense reasoning, providing a practical checkpoint-compatible route to more syntax-robust decoder-only LLMs.

Gated Tree Cross-attention for Checkpoint-Compatible Syntax Injection in Decoder-Only LLMs

TL;DR

This work tackles the brittleness of decoder-only LLMs to syntactic perturbations by introducing GTCA, a checkpoint-compatible gated tree cross-attention pathway that reads cached constituency chunk memory without altering the backbone. A token update mask and a three-stage training schedule regulate the structural updates, enabling targeted syntax improvements while preserving broad competence across tasks and backbones. Empirical results on CLOTH, MMLU, BLiMP, CoLA, HellaSwag, and Winogrande show consistent syntactic gains with minimal or no loss in MCQA performance, and layer-wise probing confirms syntax-aligned internal changes. The approach remains practical by using offline parsing, memory caching, and a gating mechanism that provides interpretability and stability, though it relies on external parsers and incurs preprocessing/compute overhead.

Abstract

Decoder-only large language models achieve strong broad performance but are brittle to minor grammatical perturbations, undermining reliability for downstream reasoning. However, directly injecting explicit syntactic structure into an existing checkpoint can interfere with its pretrained competence. We introduce a checkpoint-compatible gated tree cross-attention (GTCA) branch that reads precomputed constituency chunk memory while leaving backbone architecture unchanged. Our design uses a token update mask and staged training to control the scope and timing of structural updates. Across benchmarks and Transformer backbones, GTCA strengthens syntactic robustness beyond continued-training baselines without compromising Multiple-Choice QA performance or commonsense reasoning, providing a practical checkpoint-compatible route to more syntax-robust decoder-only LLMs.
Paper Structure (35 sections, 12 equations, 3 figures, 8 tables)

This paper contains 35 sections, 12 equations, 3 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Injecting explicit syntactic structure into pre-trained Transformers
  4. Probing syntactic representations and training-efficient adaptation
  5. Method
  6. Precomputed parse tree structure aligned to token spans
  7. Parse tree memory.
  8. MCQA-specific structure and token update mask.
  9. Hierarchical parse tree encoder for chunk memory construction
  10. Gated parse tree cross-attention injected via wrappers
  11. Three-stage training to balance syntax gains and retention
  12. Experiments
  13. Experimental Setup
  14. Model variants.
  15. Datasets and preprocessing.
  16. ...and 20 more sections

Figures (3)

  • Figure 1: Overview of GTCA (Gated Tree Cross-Attention) structural injection.Given a input (question, options, and answer field), we run an offline parser with span alignment and encode the resulting hierarchy with a parse-tree encoder to build per-layer chunk memory $C^\ell$. The chunk memory is injected into a frozen backbone layer via gated cross-attention, and the resulting structural update is used to update token states $H^\ell$.
  • Figure 2: A constituency parser produces cached span-aligned nodes $S(u)$ and update mask. chunk memory is computed by mean-pooling over $S(u)$ followed by a height-specific projection and LayerNorm.
  • Figure 3: UUAS of the Backbone, LoRA-only, Direct-Joint, and GTCA. UUAS is computed on intermediate layers, excluding the embedding layer.