Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Mitigating Position-Shift Failures in Text-Based Modular Arithmetic via Position Curriculum and Template Diversity

Nikolay Yudin

TL;DR

This work investigates why high in-distribution accuracy can coexist with catastrophic brittleness when inputs are reformatted. By studying a character-level Transformer trained to compute modular addition from text at $p=97$, the authors show a sharp position-shift robustness cliff and sensitivity to template phrasing. They propose a concise steering recipe—position diversity with a curriculum, template diversity, expression anchors, and consistency training—that substantially improves robustness to both position shifts and template OOD while preserving in-distribution performance; an ALiBi ablation fails to learn this task. The paper provides a reproducible evaluation protocol and artifacts to enable future work on steering procedural generalization under noisy supervision, with implications for reliability in instruction-following and tool-using models.

Abstract

Building on insights from the grokking literature, we study character-level Transformers trained to compute modular addition from text, and focus on robustness under input-format variation rather than only in-distribution accuracy. We identify a previously under-emphasized failure mode: models that achieve high in-distribution accuracy can fail catastrophically when the same expression is shifted to different absolute character positions ("position shift") or presented under out-of-distribution natural-language templates. Using a disjoint-pair split over all ordered pairs for p=97, we show that a baseline model reaches strong in-distribution performance yet collapses under position shift and template OOD. We then introduce a simple training recipe that combines (i) explicit expression boundary markers, (ii) position curriculum that broadens the range of absolute positions seen during training, (iii) diverse template mixtures, and (iv) consistency training across multiple variants per example. Across three seeds, this intervention substantially improves robustness to position shift and template OOD while maintaining high in-distribution accuracy, whereas an ALiBi-style ablation fails to learn the task under our setup. Our results suggest that steering procedural generalization under noisy supervision benefits from explicitly training invariances that are otherwise absent from the data distribution, and we provide a reproducible evaluation protocol and artifacts.

Mitigating Position-Shift Failures in Text-Based Modular Arithmetic via Position Curriculum and Template Diversity

TL;DR

This work investigates why high in-distribution accuracy can coexist with catastrophic brittleness when inputs are reformatted. By studying a character-level Transformer trained to compute modular addition from text at , the authors show a sharp position-shift robustness cliff and sensitivity to template phrasing. They propose a concise steering recipe—position diversity with a curriculum, template diversity, expression anchors, and consistency training—that substantially improves robustness to both position shifts and template OOD while preserving in-distribution performance; an ALiBi ablation fails to learn this task. The paper provides a reproducible evaluation protocol and artifacts to enable future work on steering procedural generalization under noisy supervision, with implications for reliability in instruction-following and tool-using models.

Abstract

Building on insights from the grokking literature, we study character-level Transformers trained to compute modular addition from text, and focus on robustness under input-format variation rather than only in-distribution accuracy. We identify a previously under-emphasized failure mode: models that achieve high in-distribution accuracy can fail catastrophically when the same expression is shifted to different absolute character positions ("position shift") or presented under out-of-distribution natural-language templates. Using a disjoint-pair split over all ordered pairs for p=97, we show that a baseline model reaches strong in-distribution performance yet collapses under position shift and template OOD. We then introduce a simple training recipe that combines (i) explicit expression boundary markers, (ii) position curriculum that broadens the range of absolute positions seen during training, (iii) diverse template mixtures, and (iv) consistency training across multiple variants per example. Across three seeds, this intervention substantially improves robustness to position shift and template OOD while maintaining high in-distribution accuracy, whereas an ALiBi-style ablation fails to learn the task under our setup. Our results suggest that steering procedural generalization under noisy supervision benefits from explicitly training invariances that are otherwise absent from the data distribution, and we provide a reproducible evaluation protocol and artifacts.
Paper Structure (55 sections, 3 equations, 4 figures, 2 tables)

This paper contains 55 sections, 3 equations, 4 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Connection to grokking.
  3. A minimal but revealing benchmark.
  4. Key observation: robustness collapses despite high in-distribution accuracy.
  5. Steering interventions.
  6. Main results.
  7. Contributions.
  8. Paper structure.
  9. Related Work
  10. Grokking on algorithmic tasks and modular arithmetic
  11. Steering or accelerating grokking-like transitions
  12. Grokking under realistic conditions and failure modes
  13. Distribution shift and domain adaptation in NLP
  14. Setup
  15. Task and dataset
  16. ...and 40 more sections

Figures (4)

  • Figure 1: Baseline vs. I1_002a on the evaluation suite (aggregated over seeds). The key gap is robustness: Eval-B and Eval-C0 improve substantially under I1.
  • Figure 2: Training curves across experiments (aggregated view). Curves illustrate that in-distribution success can coexist with robustness collapse, and that invariance training changes the outcome.
  • Figure 3: Eval-B position breakdown (aggregated over seeds). The baseline collapses under moderate shifts; position curriculum removes the cliff over the trained position range.
  • Figure 4: Final performance summary across protocols (aggregated over seeds). I1_002a improves no-anchor OOD (Eval-C0) while also performing strongly on anchor OOD (Eval-C1).