PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Robust Reinforcement Learning for Small-Scale Language Model Agents

Md Rezwanul Haque, Md. Milon Islam, Fakhri Karray

arXiv:2607.25091v1cs.AIcs.CLcs.LGmath.OCstat.CO

TL;DR

Small-language-model reinforcement learning is unstable, but its failure mechanisms and performance determinants have not been systematically established. This paper identifies reproducible PPO failures, introduces stabilization mechanisms, and finds that stable effectiveness depends on SFT fluency and reward discriminability rather than parameter count.

  • Problem

    PPO alignment of small language models lacks systematic evidence about its reproducible failure modes and stability requirements.

  • Method

    The paper evaluates a three-stage SFT, Bradley–Terry reward-modeling, and PPO system across fifteen small-model configurations with a three-layer stabilization framework.

  • Results

    Across fifteen configurations, the stabilized system remained reliable, achieved competitive instruction-tuned-baseline performance, and used significantly less training data.

  • Takeaways & Limitations

    PPO effectiveness at the SLM scale is associated with a fluent SFT prior, approximately PPL_SFT < 20, and a discriminative reward signal rather than parameter count alone.

  • Takeaways & Limitations

    Synthetic preference pairs may limit reward-model discrimination, while the 250-step PPO budget may not capture improvements from longer training.

Abstract

from arXiv · show

The alignment of Small Language Models (SLMs) in the 70--500M parameter range using reinforcement learning is often considered unstable, though the underlying failure mechanisms have not been systematically investigated. In the State-of-the-Art (SOTA) research, fifteen (model, corpus) configurations were trained using Proximal Policy Optimization (PPO). The experiments included Pythia-70M, 160M, 410M and SmolLM2-135M, 360M on the TinyStories, CNN/DailyMail, and Wikitext-103 corpora. Three reproducible failure modes were identified in small-scale language models: silent LoRA parameter freezing in standard PEFT/TRL pipelines, numerical overflow in importance ratios when using bfloat16, and catastrophic policy collapse due to reward-model error. These issues were addressed using a merge-and-reinitialize adapter technique, float32 precision during PPO updates, and a three-layer safety mechanism comprising reward whitening, importance-ratio guarding, and weight rollback. In this paper, a capacity-headroom hypothesis is proposed, which states that PPO performance at the SLM scale depends on both a fluent supervised model ($\text{PPL}<20$) and a discriminative reward signal, rather than on the number of model parameters. The proposed system converged stably in all experiments and improved preference win rate over the SFT baseline in configurations with a fluent prior and an informative reward signal. Furthermore, it outperformed instruction-tuned baselines while requiring significantly less training data. All checkpoints, preference datasets, and training scripts are publicly released$^§$.

I. INTRODUCTION · II. RELATED WORKS

This work investigates why PPO alignment is unstable for 70–500M-parameter language models, combining an empirical study of fifteen model–dataset configurations with engineering safeguards. It situates the study against large-model RLHF and SLM methods that primarily use supervised or preference learning without explicit reinforcement learning.

  • I. INTRODUCTION: SLMs in the 70–500M parameter range support resource-constrained, privacy-preserving, and on-device agents through real-time inference on common hardware.Target applications include summarization, controlled generation, and edge dialogue.
  • I. INTRODUCTION: PPO alignment remains an open problem for SLMs because learning signals are weak, policy capacity is limited, and training large models is costly.RLHF is established for large-scale models, but models below 500 million parameters have received limited application.
  • I. INTRODUCTION: The study trains five base models on three corpora through a full SFT → reward-model → PPO cycle, producing fifteen configurations with shared hyperparameters.The models are Pythia-70M, 160M, 410M, and SmolLM2-135M, 360M; the corpora are TinyStories, CNN/DailyMail, and Wikitext-103.
  • I. INTRODUCTION: Three reproducible SLM-specific PPO failures are identified: silent LoRA freezing, bfloat16 importance-ratio overflow, and reward-driven policy collapse.The reported remedies are adapter merge-and-reinitialization, float32 PPO updates, and a three-layer safety framework.
  • I. INTRODUCTION: The capacity-headroom hypothesis states that PPO effectiveness depends on a fluent SFT prior and a discriminative reward signal rather than parameter count.The hypothesis is evaluated across fifteen configurations.
  • I. INTRODUCTION: The proposed system performs comparably to public instruction-tuned baselines while using significantly less training data.Named baselines include SmolLM2-Instruct and Qwen2.5-Instruct.
  • II. RELATED WORKS: Prior RLHF work mainly targets 1.3–175 billion-parameter models, while SLM research emphasizes SFT or preference methods such as DPO, KTO, and GRPO.LoRA and QLoRA through TRL are standard parameter-efficient approaches, but gradient-flow problems and inherited stabilization techniques remain insufficiently evaluated at SLM scale.
  • II. RELATED WORKS: This work differs from DPO- and GRPO-based approaches by identifying reproducible standard-PPO failure modes at SLM scale rather than introducing a new optimization objective.The related stabilization methods include adaptive KL, GAE, and NEFTune.

III. METHODOLOGY · A. Stage 0: Data and Preference Pair Preparation

The methodology implements a three-stage RLHF system tailored to models below 500M parameters, with an agent formulation that supports structured states and multiple action types. For the empirical configurations, data are split across SFT, preference-pair construction, and evaluation, while synthetic preferences target premature stopping, incoherence, and topical mismatch.

  • III. METHODOLOGY: The RLHF pipeline comprises supervised fine-tuning, Bradley–Terry reward modeling, and PPO with a KL-divergence penalty, each tailored to models below 500M parameters.The overall data flow and three engineering mechanisms for training stability are illustrated in Fig. 1.
  • III. METHODOLOGY: The agent policy operates in a finite-horizon MDP whose state contains the prompt, generated tokens, external observations, turn index, and remaining budget.The state is represented as a structured tuple within M = (S, A, P, R, T).
  • III. METHODOLOGY: The action space includes vocabulary-token generation, external tool invocation, clarifying-query generation, and termination.These are four distinct action types available to the policy.
  • III. METHODOLOGY: The fifteen empirical configurations use a single-turn language-modeling MDP with T = 1, vocabulary actions, deterministic transitions, terminal-only rewards, and no external observations.In this form, the state reduces to the context (x, y<t).
  • A. Stage 0: Data and Preference Pair Preparation: TinyStories, CNN/DailyMail, and Wikitext-103 are separately split for SFT training, preference-pair construction, and held-out evaluation.The corpora are denoted DTS, DCNN, and DWT, respectively.
  • A. Stage 0: Data and Preference Pair Preparation: Synthetic preference pairs use truncation, sentence shuffling, and cross-example mismatch to penalize premature stopping, incoherence, and topical inconsistency.The three degradation strategies are used in equal proportion.
  • A. Stage 0: Data and Preference Pair Preparation: Truncation creates rejected continuations by cutting sentences at 50% of their length, usually mid-clause, to penalize premature stopping.The procedure specifically targets incomplete continuations.
  • A. Stage 0: Data and Preference Pair Preparation: Sentence shuffling disrupts multi-sentence order, while cross-example mismatch pairs prompts with another example’s reference text to enforce coherence and topical consistency.These corruptions target overall incoherence and prompt-reference mismatch.

B. Stage 1: Supervised Fine-Tuning

Stage 1 minimizes a causal language-modeling objective on the SFT dataset while freezing the pretrained base model and updating only LoRA adapters. The adapters are trained across attention and FFN projections with NEFTune noise, AdamW, warm-up, cosine decay, and five epochs before saving the SFT checkpoint.

  • Objective: SFT minimizes the causal language-modeling objective on dataset D_SFT = {x^(i)}_N.
  • Adapter configuration: The pretrained parameters θ_0 ∈ R^d are frozen, and updates apply only to LoRA adapters with rank r ∈ {8, 16, 32} and α = 2r.The effective weight matrix is W_0 + α/r BA.
  • Adapter configuration: LoRA is applied to attention projection matrices and dense FFN layers, with architecture-specific coverage for GPT-NeoX and Llama-style models.GPT-NeoX uses combined query-key-value and primary FFN up-projections; Llama-style models use query, key, value, gate, and downprojection matrices.
  • Optimization: NEFTune noise with α_NEFT = 5 is applied to reduce over-memorization during training.
  • Optimization: Optimization uses AdamW at a peak learning rate of 2 × 10^-5, 6% linear warm-up, cosine decay, and 5 epochs, producing the saved SFT checkpoint π_SFT.The resulting adapter Δθ_SFT is saved as the SFT checkpoint π_SFT.

C. Stage 2: Reward Model Training · D. Stage 3: PPO Fine-Tuning and Stabilization

Stage 2 trains a scalar reward model from preference data with a constrained Bradley–Terry objective. Stage 3 applies PPO with explicit KL control and addresses adapter freezing, reduced-precision ratio explosions, and distributional collapse through reinitialization and safety mechanisms.

  • C. Stage 2: Reward Model Training: The reward model replaces πSFT’s causal-language-modeling head with a linear scalar projection trained alongside a second LoRA adapter and transformer body.The scoring layer is fully fine-tuned, and training minimizes Bradley–Terry loss on preference dataset DRM.
  • C. Stage 2: Reward Model Training: Reward-model training uses AdamW for two epochs at learning rate 10^-5, batch size 8 per device, and gradient accumulation 2.The budget is limited to reduce brittle reward patterns; reward difference Δ is reported per configuration while one shared model compares models.
  • D. Stage 3: PPO Fine-Tuning and Stabilization: PPO optimizes πRL with GAE, λ = 0.95, γ = 1.0, clipping ratio ε = 0.2, and adaptive KL targeting 6.0 nats per token.The reverse KL, KL(πθ∥πSFT), penalizes divergence from the frozen SFT prior and is estimated per token from sampled roll-outs.
  • D. Stage 3: PPO Fine-Tuning and Stabilization: Merge-and-reinitialize folds the SFT LoRA adapter into base weights, then creates a fresh zero-initialized adapter for πθ and an unadapted πref.This procedure is intended to reduce PEFT gradient-flow obstruction during PPO.
  • D. Stage 3: PPO Fine-Tuning and Stabilization: PEFT gradient obstruction can leave PPO roll-outs and losses active while adapter parameters remain non-trainable, so the policy distribution does not change.The described mitigation merges the SFT adapter into the base model before initializing a fresh adapter.
  • D. Stage 3: PPO Fine-Tuning and Stabilization: In models below 200 million parameters, bfloat16 importance-ratio computation can produce ratios exceeding 10^6 during initial optimization steps.The failure arises from catastrophic cancellation in exp(log πθ − log πSFT) because bfloat16 has a seven-bit mantissa.
  • D. Stage 3: PPO Fine-Tuning and Stabilization: Long-tailed rewards and unclipped KL penalties can drive the policy toward degenerate outputs when advantage estimates exceed the clipping range.The collapse moves the policy toward regions assigned very low probability by the reference model.
  • D. Stage 3: PPO Fine-Tuning and Stabilization: The stabilization stack applies reward whitening with 3σ clipping, skips mini-batches when the importance ratio exceeds 5, and rolls back weights after NaN or Inf detection.Parameters are saved before each optimizer step and restored when invalid values appear.

IV. EXPERIMENTAL SETUP … C. Training Hyperparameters

The experimental setup spans five small language models, three corpora, and fixed training procedures across configurations. SFT, reward-model training, and PPO use specified optimization schedules and PPO stabilization parameters.

  • A. Models: Five models from two architecture families span 70M–410M parameters for analyzing architecture and parameter-count effects in RL alignment.The Pythia variants are 70M, 160M, and 410M models using GPT-NeoX BPE tokenization and Pile pretraining.
  • B. Datasets: Three corpora represent simple narratives, formal journalism, and technical content: TinyStories, CNN/DailyMail, and Wikitext-103.Each corpus contributes 10,000 examples for SFT.
  • B. Datasets: Preference pairs for reward-model training were generated through truncation, sentence shuffling, and cross-example degradation methods.The supplied passage specifies truncation as cutting chosen text mid-clause and sentence shuffling as changing sentence order.
  • C. Training Hyperparameters: All configurations used one fixed hyperparameter set without model-specific tuning, with AdamW as the optimizer.This shared setup was applied across the experimental configurations.
  • C. Training Hyperparameters: SFT used a 2×10^-5 learning rate, batch size 8×4, 6% linear warm-up, cosine decay, and 5 epochs.These settings define the supervised fine-tuning schedule.
  • C. Training Hyperparameters: Reward-model training used a 1 × 10^-5 learning rate, batch size 8 × 2, and 2 epochs.The passage reports these as the reward-model optimization settings.
  • C. Training Hyperparameters: PPO used a 5 × 10^-6 learning rate, rollout batch size 32, mini-batch size 4, and 250 optimization steps.Its configuration also used clipping range ϵ = 0.2, GAE λ = 0.95 and γ = 1.0, and adaptive KL penalty β targeting 6.0 nats.

D. Evaluation Metrics · E. Computational Environment and Reproducibility

Evaluation used 200 held-out prompts per configuration and reported perplexity, reward-based comparisons, diversity, and n-gram overlap. Experiments ran on two RTX A6000 GPUs, with scalar outputs cross-verified and an interactive verification application released.

  • D. Evaluation Metrics: Evaluation used 200 held-out prompts for each configuration, with padding tokens masked when computing perplexity on generated continuations.This masking avoided inflated perplexity scores.
  • D. Evaluation Metrics: Reward score was the mean rϕ output on model-generated responses, while reward gain ∆ was defined as PPO minus SFT reward.The signed difference is ¯rPPO − ¯rSFT.
  • D. Evaluation Metrics: Win rate estimated the probability that a PPO response received a higher score than an SFT response.
  • D. Evaluation Metrics: Lexical diversity used Distinct-1 and Distinct-2, while n-gram overlap used ROUGE-1 and ROUGE-L.
  • D. Evaluation Metrics: Absolute reward scores were compared only within configurations because Bradley–Terry loss is additively invariant; a shared reward model enabled cross-model comparisons.
  • E. Computational Environment and Reproducibility: All experiments used two NVIDIA RTX A6000 GPUs with 48 GB each and CUDA 12.1, requiring approximately 16 GPU-hours.
  • E. Computational Environment and Reproducibility: A post-hoc cross-verification script checked all 339 scalar result fields against raw evaluation outputs, and released code included an interactive verification application.

V. RESULTS ANALYSIS · A. Main Performance Analysis

Across fifteen configurations, PPO produced its strongest gains for Pythia-410M and SmolLM2-360M on TinyStories, while Pythia-70M showed little or negative change. Statistically reliable improvements occurred when the SFT prior was fluent, supporting the capacity-headroom hypothesis over raw parameter count.

  • A. Main Performance Analysis: Pythia-410M and SmolLM2-360M on TinyStories achieved the largest reward gains, with Δ=+1.355 and +0.724 and win rates of 59.9% and 59.7%.These were the clearest rightward shifts in the SFT-versus-PPO comparison.
  • A. Main Performance Analysis: Pythia-70M showed near-zero or negative changes across all domains and remained close to the SFT-PPO identity line.The identity-line position indicates little improvement over SFT.
  • A. Main Performance Analysis: Three configurations showed statistically significant improvements under a two-sided z-test on n = 200 held-out prompts with 95% confidence intervals.The reported test design evaluated whether reward differences exceeded sampling variation.
  • A. Main Performance Analysis: PPO improvement was associated with a fluent SFT prior, with gains when PPL ≲20 and no improvement when the prior was weak, rather than with raw parameter count.Table II reports PPL, reward-model scores, reward difference, and analytical win rates for SFT and PPO checkpoints.
  • A. Main Performance Analysis: Pythia-410M on TinyStories, SmolLM2-360M on TinyStories, and SmolLM2-360M on Wikitext-103 had significant gains, with p < 0.001, p < 0.001, and p < 0.05, respectively.Their confidence intervals were [+0.61, +2.10], [+0.32, +1.13], and [+0.04, +0.50].
  • A. Main Performance Analysis: The only significant regression was Pythia-410M on Wikitext-103, with Δ= −1.043 and p < 0.001.Its SFT perplexity was 25.4, above the stated fluency threshold.

SFT). REWARD SCORES ARE REPORTED ON PER-CONFIGURATION · B. Capacity-Headroom Hypothesis · C. Comparison with SOTA Instruction-Tuned Baselines

The reported results support a capacity-headroom hypothesis: PPO gains depend on SFT fluency and reward discrimination rather than parameter count alone. Domain-specific SFT and PPO-aligned checkpoints also compare favorably with instruction-tuned baselines on the tested datasets.

  • SFT). REWARD SCORES ARE REPORTED ON PER-CONFIGURATION: Per-configuration reward intervals are calculated from complete per-prompt reward distributions rather than repeated seeds.This provides a direct measure of reliability within each configuration.
  • B. Capacity-Headroom Hypothesis: Figure 3 shows a monotonic negative relationship between SFT perplexity and PPO reward gain, with an inflection point at PPL_SFT ≈20.The pattern supports the capacity-headroom hypothesis.
  • B. Capacity-Headroom Hypothesis: PPO effectiveness depends on both a fluent SFT prior and a discriminative reward signal, rather than parameter count alone.A fluent prior keeps generated samples within the reward model’s training distribution.
  • B. Capacity-Headroom Hypothesis: The PPL_SFT ≈20 boundary is an empirical decision rule for the evaluated corpora and 70–500M parameter range, not a universal threshold.Its scope is limited to the experiments described.
  • C. Comparison with SOTA Instruction-Tuned Baselines: Domain-specific SFT consistently achieves lower perplexity than SmolLM2-360M-Instruct and Qwen2.5-0.5B-Instruct on all tested datasets.This comparison is reported in Table III.
  • C. Comparison with SOTA Instruction-Tuned Baselines: +2.41 and +2.98 are the PPO-aligned checkpoint’s reward scores on TinyStories and Wikitext-103, respectively, at 360M parameters.The checkpoint achieves the highest reward scores on both datasets in the comparison.
  • C. Comparison with SOTA Instruction-Tuned Baselines: 0.310 vs. Qwen 0.282 is the Wikitext-103 lexical-diversity comparison reported for the PPO-aligned checkpoint and Qwen baseline.The result reduces most of the lexical-diversity gap with instruction-tuned baselines.

D. Text Quality and Diversity · E. Ablation: Effect of the Stabilization Mechanism

Across fifteen configurations, PPO largely preserves text diversity and reference overlap while improving diversity in selected settings. The ablation shows that only the complete stabilization mechanism enables successful training without catastrophic failure.

  • D. Text Quality and Diversity: D-1 remains within ±0.04 of the SFT prior in fourteen of fifteen configurations.The exception is SmolLM2-360M on Wikitext-103, where D-1 increases by +0.083.
  • D. Text Quality and Diversity: Table IV compares text diversity and reference-overlap metrics for the SFT prior and PPO-aligned policy across all fifteen configurations.Bold values indicate improvement in the PPO policy compared with the SFT prior.
  • D. Text Quality and Diversity: Table V qualitatively contrasts repetitive SFT continuations with PPO continuations that maintain topical relevance on TinyStories using Pythia-410M.The comparison uses prompts from the evaluation set and truncated continuations.
  • D. Text Quality and Diversity: +0.083 D-1 is achieved by SmolLM2-360M on Wikitext-103 alongside the largest reward gain in its capacity class.Within the SmolLM2 family, D-1 improves across all configurations.
  • D. Text Quality and Diversity: Pythia shows small D-1 reductions when the reward signal is uninformative, indicating movement toward the SFT distribution rather than policy collapse.ROUGE-1 and ROUGE-L remain stable across most configurations.
  • E. Ablation: Effect of the Stabilization Mechanism: Zero reward delta occurs with Naive PEFT because no gradient flows through the LoRA parameters.The ablation uses Pythia-70M on TinyStories and compares Naive PEFT, Manual unfreeze, and the proposed stabilization mechanism.
  • E. Ablation: Effect of the Stabilization Mechanism: Manual unfreeze fails during the first mini-batch because of numerical instability and the absence of a reference distribution.The proposed mechanism completes the 250-step training run without catastrophic failure.
  • E. Ablation: Effect of the Stabilization Mechanism: The stabilization mechanism combines merge-and-reinitialize, float32 precision with ρ̄ ≤5, reward whitening with 3σ clipping, and weight rollback.These mechanisms respectively restore gradient flow, prevent overflow, and prevent policy collapse.

VI. DISCUSSION · VII. CONCLUSION

The discussion presents merge-and-reinitialization, float32 updates, importance-ratio guarding, and weight rollback as engineering measures for stable SLM-scale PPO. The conclusion frames three reproducible failure modes and the resulting safety mechanism as a more reliable basis for aligning small language model agents, while identifying synthetic preference data and broader agent settings as limitations or future work.

  • VI. DISCUSSION: Merge-and-reinitialization, float32 precision, and weight rollback address key stability risks when training SLMs with TRL, LoRA, and PPO.Float32 is recommended below 200M parameters, while rollback prevents unrecoverable policy collapse.
  • VI. DISCUSSION: The convergence configuration uses an adaptive KL target of 6.0 nats and an importance-ratio cap of ¯ρ ≤5.The cap is reported as the most sensitive parameter for models with fewer than 200M parameters.
  • VI. DISCUSSION: The same configuration is applied across all fifteen cases because training remains stable across the tested learning-rate and clipping-ratio ranges.The merge operation is performed once before training with O(|θ|) complexity.
  • VI. DISCUSSION: The released framework is an unvalidated, forward-looking component, with no multi-turn experiments reported.Its release is motivated as a basis for future work on tool-call reward design and per-turn safety analysis.
  • VI. DISCUSSION: Synthetic preference pairs may limit reward-model discriminative ability, whereas human-annotated pairs could provide a stronger signal.This boundary condition defines part of the study’s scope and motivates future work.
  • VII. CONCLUSION: Three reproducible failure modes at the 70–500M scale are silent LoRA gradient freezing, bfloat16 importance-ratio overflow, and reward-driven policy collapse.The conclusion states that these issues were addressed with a three-layer cybernetic safety mechanism.
  • VII. CONCLUSION: The three-layer cybernetic safety mechanism transforms PPO updates from an unstable process into a more reliable engineering approach.The paper investigates PPO applicability for aligning small language model agents.
  • VII. CONCLUSION: The resulting checkpoints achieve competitive performance compared with the conclusion’s stated evaluation context.The supplied passage truncates immediately after this claim, so no further performance metric is available.
Loading 2607.25091v1…