Length Generalization Needs Proper Regularization

Training and evaluation code for the paper's length-generalization experiments: Mamba2 models trained from scratch or continued from a pretrained checkpoint, and SmolLM3-3B continued pretraining with sliding-window attention, the VPAD dropout family, and input-dependent attention scaling.

All released checkpoints are collected on the Hugging Face Hub: https://huggingface.co/collections/sardinelab/length-generalization-needs-proper-regularization

Contents

• Installation
• Training
  • How the experiment configs are organized
  • Dropout: the two places it can be applied
  • Configuration reference (Hydra)
• Evaluation
  • RULER (needle in a haystack)
  • HELMET
  • Using the scaled SmolLM3 checkpoints
  • Using the Mamba2 checkpoints

Installation

You need Python 3.10 or newer (tested on 3.12) and an NVIDIA GPU with a CUDA-enabled PyTorch build. The tested stack is torch 2.6 with CUDA 12.4.

git clone https://github.com/deep-spin/len-gen-regularization && cd len-gen-regularization
python -m venv .venv && source .venv/bin/activate

# 1. torch first, matching YOUR CUDA (see pytorch.org for the right index URL):
pip install "torch>=2.6" --index-url https://download.pytorch.org/whl/cu124

# 2. the project dependencies (transformers, datasets, accelerate, hydra, ...):
pip install -e .

# 3. optional: lm-eval for the RULER / short-benchmark suite
pip install -e ".[eval]"

# 4. optional but strongly recommended: fused CUDA kernels (~5-10x faster
#    Mamba2 at long context; flash attention for SmolLM3). These compile
#    against the torch you installed in step 1, so they come LAST and need
#    --no-build-isolation (plus ninja and a CUDA toolchain for flash-attn):
pip install --no-build-isolation "mamba-ssm>=2.3" "causal-conv1d>=1.5" "flash-attn>=2.7"

Everything runs from the repository root; pip install -e . only installs the dependencies, it does not install the project itself anywhere.

Training

A training run is started by naming one of the prepared experiments:

python train.py +experiment=smollm3_3b_swa_cpt_vpad

The available experiments are the files in config/hydra/experiment/ (one file per run of the paper). Checkpoints are written to outputs/<run_name>/checkpoint-<step>. For multiple GPUs, use the standard torchrun launcher:

torchrun --nproc-per-node=4 train.py +experiment=<name>

Weights & Biases. Every prepared experiment has wandb.enabled: true. Either install it (pip install -e ".[wandb]") and run wandb login (or export WANDB_MODE=offline), or turn it off for a run by adding wandb.enabled=false to the command line.

torch.compile. If you want compilation, request it on the command line (compile.torch_compile=true), not in a config file. The reason: train.py has to decide whether to keep the compiler enabled before it imports torch, and at that moment only command-line arguments are visible. A setting that exists only in a config file is therefore ignored, with a clear warning, and the run continues uncompiled.

How the experiment configs are organized

Each model family has one global config in config/hydra/model/ that holds everything the family's runs share: which base model to load, the optimizer and learning-rate schedule, the dataset, batch sizes, and so on. The files in config/hydra/experiment/ are the children: each one names its parent and adds only what makes that run different — usually the output directory, the Weights & Biases run name, and the regularizer being tested (weight decay, a dropout variant, ...). For example, experiment/smollm3_3b_swa_cpt_vpad.yaml is about a dozen lines on top of model/smollm3_3b_swa_cpt.yaml.

The four families: mamba2_1_3b_cpt_* (Mamba2-1.3B continued pretraining), mamba2_85m_lengen_pile_* (Mamba2-85M trained from scratch on the Pile), copy_jelassi_mamba2_45m_* (the synthetic copy task), and smollm3_3b_swa_cpt_* / smollm3_3b_swa_8K_* (SmolLM3-3B continued pretraining at 4K, and at 8K with attention scaling).

Dropout: the two places it can be applied

The dropout regularizers (dropout.rate + dropout.mode, implemented in dropout/) only change the model during training — at evaluation time the model behaves exactly like an unmodified one. The same perturbation (vpad or pre_linear baseline) can be inserted at two different kinds of location:

1. At the RMSNorm layers (the default): the perturbation is applied to the normalized activations that each per-layer norm produces. Because one norm feeds several projections, those projections share a single dropout mask — the query/key/value projections share one draw, and the gate/up projections share another. The model's final output norm is left alone unless you set dropout.keep_final_norm: true. This setting is used in the paper for an apples-to-apples comparison with pre-norm dropout.

2. At the linear projections (dropout.apply_to: 'q,k,v,g,u,d' for Llama-style models such as SmolLM3; i and o for Mamba2): the inputs of the selected linear layers are perturbed instead, each with its own independent mask. This also reaches places the norm-level setting cannot — for example d (the down projection) sees the output of silu(gate) * up, which never passes through a norm.

The two settings are therefore not equivalent even at the same rate and mode: they differ in which activations are perturbed and in whether masks are shared. All experiments in the paper use the RMSNorm setting (no released configuration sets apply_to); the projection-level setting is an extra capability of this repository.

Configuration for the norm-level setting, as used by the paper's dropout runs (compare config/hydra/experiment/smollm3_3b_swa_cpt_vpad.yaml):

dropout:
  rate: 0.3
  mode: vpad          # or pre_linear

and for the projection-level setting:

dropout:
  rate: 0.3
  mode: vpad
  apply_to: q,k,v,g,u   # llama/SmolLM3 letters (d = down_proj also valid);
                        # mamba2: i = mixer.in_proj

On the command line, put the comma-separated list in quotes (Hydra would otherwise read the commas as a parameter sweep): python train.py +experiment=<name> dropout.apply_to='q,k,v'.

Configuration reference (Hydra)

Configuration uses Hydra structured configs: the typed schema lives in config/schema.py (TrainingConfig, registered in config/registry.py), experiment files overlay it, and any field can be overridden on the command line with dotted names:

python train.py +experiment=mamba2_1_3b_cpt_wd                      # run a prepared experiment
python train.py +experiment=mamba2_1_3b_cpt_wd max_steps=100 optim.lr=1e-5
python train.py +experiment=mamba2_1_3b_cpt_wd --cfg job            # print resolved config, exit
python train.py +experiment=mamba2_1_3b_cpt_wd dry_run=true         # validate + print + exit

Top-level

key	meaning
model_name_or_path	HF id/path to load; empty string = build from scratch using arch (mamba2 only).
pretrained_tokenizer	tokenizer id/path (vocabulary source for from-scratch builds).
revision / tokenizer_revision	HF Hub revision pins (e.g. SmolLM3 stage3-step-4720000).
attn_implementation	flash_attention_2 (SmolLM3 training) / sdpa / eager.
output_dir, save_steps, save_total_limit, logging_steps, max_steps, seed, bf16	the usual HF Trainer settings.
resume_from_checkpoint	auto = resume from the newest checkpoint in output_dir.
dry_run	validate the config and exit before building the model.

Groups

• optim — lr, weight_decay, adam_beta1/2, adam_epsilon, max_grad_norm, optim_name (HF Trainer name, e.g. adamw_torch_fused), reference_wd_policy (the mamba_ssm convention: A_log, dt_bias, D and the norm weights are excluded from weight decay — this affects results, because decaying those state-space parameters degrades the model's dynamics).
• schedule — schedule_type (e.g. cosine_with_min_lr, warmup_stable_decay), warmup_steps, free-form kwargs (decay_ratio, min_lr_rate, ...).
• data — path (a local dataset directory, either saved with save_to_disk or a parquet export; empty = stream dataset_name from the Hub with dataset_split), block_size, shuffle_buffer_size, shuffle_seed, text_column, padded_examples (one example per padded row, used by the copy task), answer_loss_only + answer_sep, val_sets / val_batch_size / val_max_samples (the in-training generation evaluation), eval_metric (exact_match_accuracy, the default — whole-sequence exact match | token_accuracy — position-wise matches of the free-running generation | teacher_forced_accuracy — the "character accuracy" of "Understanding and Improving Length Generalization in Recurrent Models", Buitrago and Gu (2025): the per-position argmax over one forward pass of the gold sequence), use_doc_masking, no_doc_mixing, append_eos_between_docs.
• arch — the from-scratch Mamba2 architecture (hidden_size, expand, num_heads, head_dim, state_size, num_hidden_layers, tie_word_embeddings, mamba2_reference_init). Ignored when model_name_or_path is set.
• swa — the SmolLM3 sliding-window setup (enabled, window, rope_theta): a config-only change applied when the model is loaded; layers with rotary embeddings become sliding_attention, the native NoPE layers stay full_attention.
• dropout — the dropout family described above: rate, mode (vpad, the default, or pre_linear), apply_to (projection letters: llama q,k,v,g,u,d; mamba2 i,o), layers (restrict to a subset of blocks, e.g. '0,3,7-11'), keep_final_norm, skip_patterns / skip_module_names, include_mixer_norms, no_detach_beta, no_detach_alpha. Alias: mode: pre_norm (or residual) redirects the whole block to the residual_dropout group at load time (rate and layers carry over).
• attention_dropout — the native HF dropout on attention probabilities: rate, layers. Applied to the config before the model is loaded.
• attention_scaling — input-dependent per-head query scaling, beta_h(p, x) = 1 + coef_h(x) * log2(p + 2) ** pow_h(x): enabled, coef_bound, pow_bound, identity_init, proj_bias, coef_activation (sigmoid | softplus_shifted), weight_scale_alpha. The settings are recorded in each checkpoint's config.json as attention_scaling_* entries, and the evaluation loaders rebuild the scaling from those entries.
• wandb — enabled, project, run_name.
• compile — torch_compile (per-decoder-layer compilation; must be given on the command line, see the note above), backend, disable_dynamo.
• eval_extra — the per-position cross-entropy evaluation: eval_only (evaluate a checkpoint without training), dataset_file, output_name, text_keypaths.

Evaluation

The released checkpoints are collected here: https://huggingface.co/collections/sardinelab/length-generalization-needs-proper-regularization. The evaluation entry points are eval/lm_eval_patched.py (RULER and the short benchmarks, through lm-eval), eval/helmet_eval.py (the HELMET long-context suite), and eval/eval_synth_em.py (exact-match generation on the synthetic copy task). All of them accept either a local checkpoint directory or a Hub repository id.

RULER (needle in a haystack)

The RULER tasks hide a small piece of information (the "needle") inside a long filler text and ask the model to retrieve it. They run through eval/lm_eval_patched.py, which is a drop-in replacement for the lm-eval command that additionally restores this project's model modifications when loading a checkpoint. Scaled Hub checkpoints need trust_remote_code=True in model_args; the plain dropout checkpoints do not.

Two length settings matter: --metadata max_seq_lengths chooses how long the generated test inputs are, and max_length in model_args must be larger than that (with some room for the generated answer) — otherwise lm-eval silently cuts off the beginning of the input and the long-context score is meaningless:

python eval/lm_eval_patched.py \
    --model hf \
    --model_args pretrained=vapavlo/SmolLM3-SWA-8K-VPAD-p30-scaled,trust_remote_code=True,dtype=bfloat16,max_length=32768 \
    --tasks niah_single_1 \
    --metadata '{"max_seq_lengths":[16384]}' \
    --batch_size 1 \
    --output_path eval_results/hub_scaled_niah

HELMET

HELMET is a long-context benchmark suite from Princeton NLP. eval/helmet_eval.py runs it with this project's checkpoint loading in place. Clone HELMET once and install its two scoring dependencies; the ICL task (many-shot in-context classification with 128K-token prompts) needs no dataset download — its five datasets are fetched from the Hub automatically:

git clone https://github.com/princeton-nlp/HELMET
pip install pytrec_eval rouge_score

python eval/helmet_eval.py --helmet-src HELMET \
    --config HELMET/configs/icl.yaml \
    --model_name_or_path vapavlo/SmolLM3-SWA-8K-VPAD-p30-scaled \
    --output_dir eval_results/helmet_icl_scaled \
    --no_torch_compile

A Hub repository id is downloaded to a local copy automatically. Results are written as one .json / .json.score file pair per dataset under --output_dir; add --max_test_samples 50 for a quicker pass. The other HELMET task configs (recall, rag, longqa, ...) work the same way but require HELMET's data download — see their repository.

ICL at 256K and 512K. HELMET's own icl.yaml goes up to 128K. The prompt length is controlled by the number of demonstrations encoded in the dataset names, so this repository ships two pre-scaled configs: eval/helmet_configs/icl_256k.yaml (twice the demonstrations of the 128K suite) and icl_512k.yaml (demonstration counts measured for the SmolLM3 tokenizer, with a slight surplus so that HELMET's own truncation cuts the input to exactly the 524288-token budget; demonstrations repeat once a label's pool of unique examples is used up, as in HELMET's own sampler). Same command, different config:

python eval/helmet_eval.py --helmet-src HELMET \
    --config eval/helmet_configs/icl_256k.yaml \
    --model_name_or_path vapavlo/SmolLM3-SWA-8K-VPAD-p30-scaled \
    --output_dir eval_results/helmet_icl_256k --no_torch_compile

Measured prompt lengths: about 212K tokens for the 256K config (the same fill ratio as the stock 128K config) and exactly the 524288-token budget for the 512K config.

Using the scaled SmolLM3 checkpoints

Checkpoints trained with input-dependent attention scaling carry extra per-layer weights for the scaling and come bundled with a small custom modeling file; loading them therefore needs trust_remote_code=True (see trust_remote_code/README.md for details). All other SmolLM3 checkpoints (the dropout runs) are stock SmolLM3ForCausalLM models and load like any Hugging Face model.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "vapavlo/SmolLM3-SWA-8K-VPAD-p30-scaled"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
    repo, trust_remote_code=True, dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",   # sdpa / eager work too
).cuda().eval()

ids = tok("The main driver of climate change is", return_tensors="pt").input_ids.cuda()
out = model.generate(ids, max_new_tokens=32, do_sample=False)
print(tok.decode(out[0], skip_special_tokens=True))

Using the Mamba2 checkpoints

The released Mamba2 checkpoints are standard Hugging Face Mamba2ForCausalLM models (config, model.safetensors, and the GPT-NeoX tokenizer bundled). They load with stock transformers — no code from this repository is required. Tested with transformers==5.8.0, torch==2.6.

The complete flow — works with or without the fused kernels installed:

# 1. Load. dtype="auto" keeps the stored per-tensor dtypes (bf16 weights,
#    fp32 SSM scalars — see "Precision" below).
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "vapavlo/Mamba2-1.3B-VPAD-p20"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, dtype="auto").cuda().eval()

# 2. Confirm which path you are on (True = fused kernels active; the flag
#    is set when the first Mamba2 layer is built, so read it AFTER loading):
import transformers.models.mamba2.modeling_mamba2 as m2
print("fused kernels:", m2.is_fast_path_available)

# 3. Generate.
ids = tok("The main driver of climate change is", return_tensors="pt").input_ids.cuda()
out = model.generate(ids, max_new_tokens=32, do_sample=False)
print(tok.decode(out[0], skip_special_tokens=True))

Precision — important. The checkpoints store most weights in bfloat16 but keep the SSM scalar parameters (A_log, D, dt_bias) in float32 on purpose: those parameters are precision-critical for the state-space dynamics.

• dtype="auto" (recommended) keeps the stored per-tensor dtypes.

• The default (float32) upcasts everything — safe, at twice the memory.

• Do not pass dtype=torch.bfloat16: it converts A_log/D/dt_bias to bfloat16 while loading, and the lost precision cannot be recovered afterwards. If you must force bf16 storage, extend the class-level float32 list first:

  from transformers.models.mamba2 import Mamba2PreTrainedModel
  Mamba2PreTrainedModel._keep_in_fp32_modules_strict = \
      list(Mamba2PreTrainedModel._keep_in_fp32_modules_strict or []) + \
      ["A_log", "D", "dt_bias"]