Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Sample-Efficient Language Modeling with Linear Attention and Lightweight Enhancements

Patrick Haller, Jonas Golde, Alan Akbik

TL;DR

The paper tackles the challenge of sample-efficient language modeling under BabyLM’s strict data and epoch limits by swapping self-attention for a linear-time mLSTM token mixer (BLaLM) and applying lightweight enhancements such as ShortConv, Sliding Window Attention, Dynamic Modulation, and Hedgehog feature maps. It also investigates the Muon optimizer as a superior alternative to AdamW for matrix-shaped parameters, all within a curated, pedagogy-focused corpus. Across experiments, linear attention with local enhancements yields strong zero-shot performance in low-resource settings and competitive results at higher scales, while Muon stabilizes training and reduces perplexity. Collectively, the work offers practical, data-efficient design choices for compact language models that avoid reliance on scale alone and provides actionable guidance for educational NLP applications within tight resource constraints.

Abstract

We study architectural and optimization techniques for sample-efficient language modeling under the constraints of the BabyLM 2025 shared task. Our model, BLaLM, replaces self-attention with a linear-time mLSTM token mixer and explores lightweight enhancements, including short convolutions, sliding window attention with dynamic modulation, and Hedgehog feature maps. To support training in low-resource settings, we curate a high-quality corpus emphasizing readability and pedagogical structure. Experiments across both STRICT and STRICT-SMALL tracks show that (1) linear attention combined with sliding window attention consistently improves zero-shot performance, and (2) the Muon optimizer stabilizes convergence and reduces perplexity over AdamW. These results highlight effective strategies for efficient language modeling without relying on scale.

Sample-Efficient Language Modeling with Linear Attention and Lightweight Enhancements

TL;DR

The paper tackles the challenge of sample-efficient language modeling under BabyLM’s strict data and epoch limits by swapping self-attention for a linear-time mLSTM token mixer (BLaLM) and applying lightweight enhancements such as ShortConv, Sliding Window Attention, Dynamic Modulation, and Hedgehog feature maps. It also investigates the Muon optimizer as a superior alternative to AdamW for matrix-shaped parameters, all within a curated, pedagogy-focused corpus. Across experiments, linear attention with local enhancements yields strong zero-shot performance in low-resource settings and competitive results at higher scales, while Muon stabilizes training and reduces perplexity. Collectively, the work offers practical, data-efficient design choices for compact language models that avoid reliance on scale alone and provides actionable guidance for educational NLP applications within tight resource constraints.

Abstract

We study architectural and optimization techniques for sample-efficient language modeling under the constraints of the BabyLM 2025 shared task. Our model, BLaLM, replaces self-attention with a linear-time mLSTM token mixer and explores lightweight enhancements, including short convolutions, sliding window attention with dynamic modulation, and Hedgehog feature maps. To support training in low-resource settings, we curate a high-quality corpus emphasizing readability and pedagogical structure. Experiments across both STRICT and STRICT-SMALL tracks show that (1) linear attention combined with sliding window attention consistently improves zero-shot performance, and (2) the Muon optimizer stabilizes convergence and reduces perplexity over AdamW. These results highlight effective strategies for efficient language modeling without relying on scale.

Paper Structure

This paper contains 43 sections, 4 equations, 3 figures, 15 tables.

Table of Contents

  1. Introduction
  2. Preliminaries and Related Work
  3. Data Curation
  4. Data Sources
  5. Filtering Pipeline
  6. FineWeb-Edu
  7. Gutenberg Fiction
  8. TinyStories
  9. Simple Wikipedia
  10. Cosmopedia (WikiHow & Math)
  11. CHILDES (Child-Directed Speech)
  12. Model Architecture and Optimization
  13. Architectural Enhancements
  14. Training Setup
  15. Experiments
  16. ...and 28 more sections

Figures (3)

  • Figure 1: Overview of the BLaLM architecture. The standard self-attention module is replaced by an mLSTM token mixer. Optional enhancements such as sliding window attention (SWA) can be integrated and combined with mLSTM outputs.
  • Figure 2: LLM-based prompt used to assign a custom educational scores to FineWeb-Edu samples. The prompt includes a 5-point additive scoring rubric focusing on pedagogical value, readability, and coherence.
  • Figure 3: Layer-wise development of dynamic modulation weights ($\alpha$) during training for the bounded DynMod variant. We apply the tanh function to stabilize values. Later layers show increased reliance on local mixing.