Source-linked AI summary

Evaluating Tiny Recursive Models Across Training for Code Generation

Anjani Sirivella, Aanisha Newaz, Glaucia Melo

arXiv:2608.29376v1cs.AIcs.LGcs.SE

TL;DR

Small recursive code models are usually assessed with teacher-forced fit or single-checkpoint accuracy, leaving their free-running behavior and training-trajectory rankings unclear. This paper compares TRM-AR with parameter- and depth-matched transformers across training and seeds, finding parameter-efficiency gains but a depth-matched transformer optimum and substantial compute cost. The results support trajectory-aware joint evaluation of fit and generation, while remaining bounded by the study’s synthetic setting and unresolved recursion-versus-supervision attribution.

  • Problem

    Existing recursive-model evaluations rarely establish whether teacher-forced advantages survive free-running natural-language-to-code generation across training.

  • Method

    The study compares TRM-AR with parameter- and depth-matched transformer controls on natural-language-to-Python generation, tracking fit and generation across training and three seeds.

  • Results

    At equal parameters, TRM-AR fits, generates, and generalizes better; at equal effective depth, the larger plain transformer fits and generates better at its validation optimum.

  • Takeaways & Limitations

    Recursive models should be evaluated jointly on fit and free-running generation across the training trajectory rather than at a single checkpoint.

  • Takeaways & Limitations

    TRM-AR combines recursion with deep supervision, so the equal-parameter findings credit the configuration as a whole rather than isolating recursion’s mechanism.

Abstract

from arXiv · show

Code generation increasingly relies on large transformer models, whose capability advances with scale. Yet such a scale is costly, creating demand for small models, especially where data is limited. Recursive models address this by reusing a single block to add depth rather than stacking independent layers. Such models are typically evaluated by teacher-forced fit (next-token loss on ground-truth prefixes) or task accuracy, at a single checkpoint, whereas code is produced by free-running generation, where the model extends its own output. Whether a teacher-forced advantage survives free-running generation, and whether it holds across training, remains open. To study both, we compare a ~28M-parameter autoregressive Tiny Recursive Model (TRM-AR) on natural-language-to-Python code generation against parameter-matched and depth-matched controls, tracking fit and generation across 40 epochs and three seeds. The fit ranking between the recursive model and the depth-matched control reverses twice. Selecting each checkpoint by validation loss and examining the trajectory yields a consistent comparison. At equal parameters, TRM-AR fits, generates, and generalizes better than the parameter-matched control while recovering approximately 45% of the validation-loss gap and 57% of the generation-quality gap between the two controls, at roughly 175 times the per-step cost of the parameter-matched control. However, at equal effective depth, the larger transformer fits and generates better at its validation optimum, suggesting TRM-AR's advantage lies in resistance to overfitting, not greater capability. These findings suggest that recursive code generation models should be evaluated jointly on fit and generation across the training trajectory rather than at a single checkpoint.

I. INTRODUCTION

The paper studies whether recursive models can generate code effectively at small scale, where large transformers are costly and existing evaluations rarely test free-running generation across training. It compares TRM-AR with matched transformer controls and finds that rankings depend on both evaluation regime and checkpoint selection.

  • Large code models are costly to train and serve, motivating small models that generate code well under limited-data conditions.
  • Recursive models add computational depth by reusing one small weight-shared block, but prior evaluations primarily use teacher-forced fit or single-checkpoint task accuracy.
  • The study tracks teacher-forced fit and free-running generation across training for TRM-AR, parameter-matched, and depth-matched transformer controls.
  • The fit ranking between TRM-AR and the depth-matched control reverses twice, with the later reversal arising from overfitting as the control worsens while TRM-AR holds steady.
  • The paper recommends reporting fit and free-running generation jointly across the training trajectory rather than relying on a single checkpoint.
  • At equal parameters, TRM-AR improves fit, judged generation, and generalization, whereas at equal effective depth the larger transformer performs better at its optimum.
  • Recursion is parameter-efficient but not compute-efficient, costing roughly 175 times the parameter-matched control’s per-step computation.

II. RELATED WORK

The paper situates its study at the intersection of recursive architectures and evaluation sensitivity, combining weight-shared models with checkpoint- and metric-aware comparison.

  • The work draws on recursive, weight-shared architectures as model foundations and evaluation research on metric and checkpoint sensitivity as comparison guidance.

A. Recursive and Weight-Shared Transformers

Prior work shows that shared blocks can obtain depth with fewer parameters, but existing evidence largely uses converged checkpoints and does not establish whether recursive advantages transfer to free-running code generation across training.

  • Recursive architectures obtain depth by repeatedly reusing a shared block instead of stacking independent layers.
  • TRM is a minimal two-layer recursive model that adds deep supervision over iterative latent refinement, beyond weight sharing alone.
  • Prior small-model studies mainly report teacher-forced fit or puzzle accuracy at converged checkpoints, leaving free-running natural-language-to-code generation across training untested.
  • Matched controls are necessary because recursive and plain transformers differ simultaneously in parameter count and effective depth.
  • This study retains deep supervision, tracks fit and free-running generation across the full trajectory, and matches parameters and effective depth while measuring rather than matching per-step cost.

C. Metric and Checkpoint Sensitivity in Evaluation

The evaluation framework addresses train–inference mismatch and checkpoint sensitivity by combining matched model comparisons, trajectory tracking, and proxy-based generation metrics on a controlled code corpus.

  • Teacher-forced next-token loss can diverge from free-running generation, so held-out loss alone does not establish generated-code quality.
  • The paper’s contribution is to chart where recursive and matched controls cross, plateau, or diverge across training and seeds.
  • The study compares TRM-AR with transformers matched on parameter count or effective depth.
  • The corpus is a shared synthetic collection of short natural-language instruction–code pairs, restricted to Python records.
  • Validation selects checkpoints while test data is reserved for final evaluation; exact prompt overlap across splits is zero, but fuzzy or semantic overlap is not measured.
  • The synthetic, templated corpus is a controlled architecture testbed rather than a benchmark establishing absolute quality or transfer to human-authored code.

B. Models

The study compares three models built from an identical transformer block, varying whether depth comes from recursive reuse or independent layer stacking. TRM-AR uses recursive computation and deep supervision, while controls match its parameter budget or effective depth.

  • Shared backbone: All three arms share an identical transformer block, isolating how computational depth is realized as the main architectural difference.The shared backbone uses decoder-only transformer layers; recursion versus stacking remains, while deep supervision is not independently separated.
  • TRM-AR: TRM-AR reaches effective depth 20 by recursively reusing a shared two-layer block at approximately a two-layer model’s parameter budget.Its recursive computation maintains latent and answer states and applies the shared block repeatedly per forward pass.
  • Matched controls: vanilla-2L matches TRM-AR’s parameter budget with two independent layers, differing by fewer than 0.003% of parameters.The small mismatch consists of TRM-AR’s learned initial states and inactive halting head.
  • Matched controls: vanilla-20L matches TRM-AR’s effective depth with twenty independent layers and approximately 10× its non-embedding parameters.The depth-matched control has 1.68× the total parameter count.
  • Recursive training and inference: Training uses three-step deep supervision, while each forward pass recomputes recursive states and executes the full fixed recursion without adaptive halting.Only the final recursion cycle receives gradients within a step, and no key/value cache is maintained during generation.

C. Training

All arms use the same optimization protocol and a fixed 40-epoch budget, with snapshots spanning early training through overfitting. Checkpoints are selected post hoc by validation loss, with three seeds per arm.

  • Optimization: All three arms use an identical optimization protocol, making matched-epoch comparisons primarily architectural.The protocol uses AdamW, gradient clipping, batch size 48, and fp32 training.
  • Optimization: The learning rate warms up for 2,000 steps and then remains constant, while EMA weights are used for validation, generation, and checkpointing.This flat-after-warmup schedule keeps snapshots in the same optimization regime rather than placing them at different decay stages.
  • Objective: Training applies completion-only cross-entropy, excluding instruction and padding tokens; TRM-AR averages the loss over three deep-supervision steps.The vanilla controls use the corresponding single-pass objective.
  • Training budget and snapshots: 40 epochs cover rapid improvement, plateau, and overfitting, with snapshots at nine specified epochs that are denser early in training.The deepest arm’s pilot validation loss reached its minimum before the full budget and then rose.
  • Checkpoint selection: Each arm’s reported checkpoint is its lowest-validation-loss epoch, although TRM-AR and vanilla-2L remain right-censored at epoch 40.Three seeds—42, 65, and 108—are used in the repeated-data regime.

D. Evaluation

Every checkpoint is evaluated using both teacher-forced fit and free-running generation. The evaluation combines perplexity, surface generation proxies, and LLM judging across held-out prompts and sampled completions.

  • Evaluation design: Every checkpoint is scored identically across arms and snapshots using teacher-forced fit and free-running generation.This directly evaluates whether fit and generation agree across training.
  • Teacher-forced fit: Teacher-forced fit is measured by test-set perplexity from completion-only next-token cross-entropy, using TRM-AR’s deployed single-pass inference path.The shared tokenizer, vocabulary, split, and masking make token-level perplexity comparable across arms.
  • Free-running generation: Free-running evaluation generates completions for held-out prompts and computes AST-parse rate and BLEU-4 as surface proxies.Code extraction checks a fenced Python block, then a bare code fence, then raw generated text.
  • Free-running generation: BLEU-4 is reported for completeness but not used to rank arms because it weakly correlates with functional code correctness.No execution-based correctness metric is reported because the dataset has no associated unit tests.
  • Free-running generation: An LLM judge provides a complementary generation-quality signal alongside the surface proxies.Completions are rated on syntax, logic, and other rubric-defined axes using integer scores from 0 to 10.

E. Reproducibility

The study documents its environment, corpus construction, randomness, and computational costs to support reproducible interpretation. Reproducibility is assessed across seeds rather than claimed bit-for-bit for individual runs.

  • Environment and data: The environment is pinned to 129 package versions, Python 3.10.14, PyTorch 2.6.0, CUDA 12.1, and a single NVIDIA A100 80 GB GPU.The corpus is also identified by a SHA-256 fingerprint over its partitions and build manifest.
  • Determinism: Random seeds control initialization, data ordering, and sampling, but nondeterministic CUDA operations prevent bit-for-bit reproducibility of individual runs.The study instead uses three seeds per arm and interprets effects through across-seed consistency.
  • Compute accounting: The arms are matched on parameters or effective depth, not compute, so recursive per-step cost is measured rather than controlled.TRM-AR reaches effective depth 20 at roughly the vanilla-2L parameter budget by repeatedly applying the shared block.
  • Compute accounting: Measured on an A100, TRM-AR takes approximately 3.0 seconds per optimizer step, versus 1.25 seconds for its iso-parameter twin and about 2.4× the iso-depth model.The study reports wall-clock time because parameter-based FLOP estimates understate sequential recursive cost.

IV. RESULTS

Validation-selected checkpoints establish the stable ordering, while the full trajectory reveals crossovers, overfitting, and fit–generation decoupling that single snapshots can misrepresent.

  • At validation-selected checkpoints, vanilla-20L leads both fit and judged generation, followed by TRM-AR and vanilla-2L.
  • Iso-parameter comparison: At equal parameters, TRM-AR leads vanilla-2L on fit across all nine epochs, with seed-consistent differences, while generation leads from epoch 4 onward.
  • Iso-depth comparison: At equal effective depth, the fit ranking reverses twice: vanilla-20L leads from epoch 8 through 28, then TRM-AR leads at epoch 40 because vanilla-20L overfits.
  • Generalization and overfitting: At epoch 40, the generalization gaps are 0.13 for TRM-AR, 0.27 for vanilla-2L, and 1.04 for vanilla-20L.
  • Fit–generation decoupling: From epoch 20 to 40, vanilla-20L’s validation loss worsens from 1.15 to 1.45 while its judge score declines only from 3.12 to 3.01.
  • Single-snapshot artifacts: Single-checkpoint readings can falsely declare either vanilla-20L or TRM-AR superior, whereas validation selection and the full trajectory distinguish overfitting from genuine improvement.

V. DISCUSSION

The discussion interprets the results as time-indexed and evaluation-dependent: recursive models offer parameter-budget advantages and more stable generalization, but at substantial sequential-computation cost.

  • The comparison is time-indexed: fit and generation can rank the three arms differently, and the fit ranking reverses twice across training.
  • Realizing effective depth 20 by looping a two-layer block costs roughly 175× the iso-parameter control’s per-step time and about 2.4× the iso-depth control’s.
  • At equal parameters, recursion improves free-running generation from epoch 4 onward, while BLEU-4 shows no stable difference.
  • TRM-AR’s smallest generalization gap is consistent with parameter sharing acting as an implicit regularizer, but the effect is entangled with depth, deep supervision, and architecture.
  • The study extends recursive-efficiency evidence to free-running code generation across training, but its comparisons differ from concurrent compute-matched studies in task, matching criterion, and configuration.
  • For small code models under fixed parameter budgets and limited data, the results favor recursion if its roughly 175× per-step compute cost is acceptable.

VI. LIMITATIONS

The study's conclusions are bounded by confounds in the TRM-AR configuration, proxy-based evaluation, synthetic data, repeated-data training, limited seeds, and a fixed 40-epoch budget.

  • Internal validity: TRM-AR's equal-parameter result credits recursion and deep supervision together because the study does not ablate them.An independent HRM analysis identified deep supervision rather than recursion as the apparent primary driver of that model's gains.
  • Construct validity: Execution-based correctness is unavailable, so generation quality relies on AST-parse rate, BLEU-4, and an uncalibrated LLM judge.The judge is read comparatively between arms at matched training points; BLEU correlates only weakly with functional correctness.
  • External validity: The synthetic, templated, LLM-generated corpus differs from human-authored code, leaving proxy levels and transfer to naturally occurring code unestablished.Across-training conclusions also apply specifically to repeated passes over a fixed, data-constrained corpus.
  • Conclusion validity: Three seeds are too few for statistical significance, so the study reports effect directions, raw magnitudes, and sign-consistency instead of p-values or standardized effect sizes.This bounds the strength of statistical claims while preserving the reported directional comparisons.
  • Conclusion validity: The fixed 40-epoch budget right-censors the two parameter-lean arms, whose validation loss is still decreasing at epoch 40.The iso-depth ordering is therefore claimed only within this training budget.
  • Scope: Within these constraints, the study provides a controlled, seed-replicated characterization of teacher-forced fit and free-running generation across training.This describes the study's supported empirical scope rather than absolute code-generation ability.

VII. CONCLUSION

The study shows that architecture rankings can change substantially across training, so evaluating recursive code-generation models over full trajectories is more reliable than relying on a single checkpoint. It also identifies unresolved questions about disentangling deep supervision, generalizing beyond the synthetic setting, and testing other recursive architectures.

  • VII. CONCLUSION: Architecture rankings changed substantially across training, making trajectory-aware evaluation more reliable than single-checkpoint comparisons in similar settings.The authors advocate trajectory-based evaluation for recursive models and suggest broader value for emerging small-scale architectures with different learning dynamics.
  • Future work: A deep-supervised shallow transformer is needed to disentangle recursive parameter sharing from deep supervision.This is the first proposed direction for clarifying the mechanism behind the observed comparison.
  • Future work: Larger models, more diverse code-generation benchmarks, and execution-based metrics such as pass@k are needed to establish the generality of the crossover behaviour.These extensions would test whether the findings transfer beyond the present corpus and proxy metrics.
  • Future work: Additional recursive and weight-sharing architectures could determine whether trajectory-dependent rankings generalize beyond the TRM studied here.The proposed comparison distinguishes a broader property of recursive models from a TRM-specific effect.
Loading 2608.29376v1…