LaCy: What Small Language Models Can and Should Learn is Not Just a Question of Loss

Abstract: LLMs have consistently grown to compress more world knowledge into their parameters, but the knowledge that can be pretrained into them is upper-bounded by their parameter size. Especially the capacity of Small LLMs (SLMs) is limited, leading to factually incorrect generations. This problem is often mitigated by giving the SLM access to an outside source: the ability to query a larger model, documents, or a database. Under this setting, we study the fundamental question of \emph{which tokens an SLM can and should learn} during pretraining, versus \emph{which ones it should delegate} via a \texttt{<CALL>} token. We find that this is not simply a question of loss: although the loss is predictive of whether a predicted token mismatches the ground-truth, some tokens are \emph{acceptable} in that they are truthful alternative continuations of a pretraining document, and should not trigger a \texttt{<CALL>} even if their loss is high. We find that a spaCy grammar parser can help augment the loss signal to decide which tokens the SLM should learn to delegate to prevent factual errors and which are safe to learn and predict even under high losses. We propose LaCy, a novel pretraining method based on this token selection philosophy. Our experiments demonstrate that LaCy models successfully learn which tokens to predict and where to delegate for help. This results in higher FactScores when generating in a cascade with a bigger model and outperforms Rho or LLM-judge trained SLMs, while being simpler and cheaper.

• The paper demonstrates that loss signals alone are insufficient for SLM token selection, integrating spaCy-based factual annotation to guide delegation.
• It introduces a dual pipeline that triggers <CALL> on high-loss factual tokens, achieving improved FactScore and reduced fact leakage in biographical generation.
• Empirical findings on GPT-2 variants show that delegated factual predictions preserve SLM NLU capabilities while mitigating risky fact memorization.

Authoritative Summary of "LaCy: What Small LLMs Can and Should Learn is Not Just a Question of Loss" (2602.12005)

Motivation and Problem Statement

The paper investigates the learnability boundaries of Small LLMs (SLMs), emphasizing their constrained parametric memory and susceptibility to factual errors. While large LLMs saturate benchmarks by memorizing world knowledge, SLMs—defined by parameter counts typically below 1B—must balance efficient token prediction against their limited capacity. This operational constraint is often mitigated by cascade mechanisms: SLMs predict what they can, then delegate to a larger partner (M, DB, or LLM) via a special <CALL> token. The authors pose the fundamental question: which tokens should be "learned" by an SLM, and which should trigger a <CALL> deferral?

Conventional pretraining minimizes token-wise cross-entropy loss. However, the paper demonstrates that loss signals are inadequate for token selection since they are agnostic with respect to factuality versus plausibly acceptable continuations. Many tokens may have high loss due to non-exact matches but are still acceptable (non-factual, e.g., synonyms or grammatical variations), whereas factual tokens (entities, dates) require precise prediction to avoid hallucinations.

LaCy Framework and Methodological Innovations

LaCy is introduced as a token-selection framework that augments the loss signal with spaCy-based grammatical and factual annotation to distinguish between tokens that SLMs can safely learn and those that should be delegated. The method identifies factual tokens (e.g., new entities, dates) with spaCy, then delegates those with high loss. Non-factual tokens may have high loss without semantic risk, so they are not automatically delegated. Formally, within each batch, LaCy replaces ground-truth targets with <CALL> tokens where the spaCy fact label and loss cross a threshold (top $n$\% highest-loss factual tokens).

(Figure 1)

Figure 1: LaCy pipeline - during pretraining, SLM's prediction and loss are parsed for factuality; factual tokens with high loss trigger training towards <CALL> output, which at inference delegates to a larger cascade model.

At inference, the SLM generates tokens autoregressively, emitting <CALL> when its capacity is insufficient for factual correctness. This triggers the cascade partner to produce the next token, supporting high factuality in downstream outputs while minimizing factual memorization in the SLM itself.

Empirical Evaluation and Analysis

Factual Accuracy and Targeted Delegation

The main experimental setup involves pretraining GPT-2 architectures (334M and 1.3B) on the dwiki dataset. Biographical generation is evaluated with FactScore, which quantifies the proportion of generated atomic facts supported by the true Wikipedia entry. LaCy achieves the highest FactScore among tested methods, notably outperforming Rho-based methods and LLM judge annotation baselines—even with a relatively limited cascade partner (Llama 3.2 1B).

Figure 2: LaCy achieves superior FactScore while minimizing fact leakage by successfully learning factual delegation locations and offloading factual predictions.

Ablations demonstrate that spaCy-only annotation (without loss selection) underperforms LaCy, confirming the necessity of loss-augmented factual token delegation. The use of a reference model’s loss (Rho-1) yields marginal improvements, but incurs significant computational overhead.

Fact Leakage and Parametric Memory

To assess factual memorization, the SLM is evaluated in isolation (with <CALL> suppressed) on QA tasks. LaCy leads to the lowest fact leakage, confirming that its delegation results in minimal factual entanglement in the SLM’s parametric memory.

Figure 2: Left - FactScore for biography generation with cascade partner; Right - fact leakage in direct SLM generation (no cascade). LaCy optimally balances factual precision and parameter memory.

Token-Type Sensitivity: Loss versus Acceptability

The authors empirically validate that loss is not correlated with downstream factual performance (FactScore) in cascaded setups. High loss may occur for non-factual tokens (e.g., grammatical alternatives), which do not risk hallucination. Only the combination of spaCy factual annotation and loss reveals tokens where SLM predictions are likely unacceptable, motivating targeted delegation.

Figure 3: Difference between accuracy (exact matching) and acceptability; loss is blind to token type, but spaCy parsing differentiates factual tokens needing deferral from non-factual tokens.

Validation loss metrics (call, non-call, and total) are shown to be poor predictors for factual accuracy in cascade scenarios, diverging from historical findings in scaling law literature [kaplan2020scaling].

NLU Performance and Model Capacity

NLU tasks (ARC-Easy, HellaSwag, PIQA, SIQA) are employed to measure side effects. Offloading facts does not impair SLM NLU skills, nor does aggressive offloading (in ablations) improve NLU, establishing that factual memorization and natural language comprehension are largely orthogonal in SLM training.

Throughput and Practical Considerations

LaCy’s integration with spaCy on purely CPU hardware enables large-scale pretraining without the labeling overhead associated with LLM judge or reference loss methods, which require substantial GPU time. The pipeline is amenable to online labeling during data loading.

Implications, Limitations, and Future Directions

The findings contradict the idea that loss-based methods alone suffice for optimal SLM training. SLMs should be trained to only memorize what is safe for their capacity—non-factual, high-loss tokens with semantic acceptability—and delegate on factual, high-loss tokens. The LaCy framework, conceptually simple, offers measurable gains in factual accuracy in cascaded settings and purges fact memorization from SLMs, with minimal throughput cost.

This approach has ramifications for cost-efficient LLM deployment, LLM cascading architectures [gupta2024language], fact hallucination mitigation, and downstream parameter-efficient fine-tuning. Future work should focus on robust scaling, joint architectures with multi-modal or large KB cascades, and refined token-selection heuristics that combine other signals (contextual uncertainty, knowledge graphs).

Conclusion

LaCy establishes that optimal token selection for SLM pretraining is not a function of loss alone; factual token identification and targeted delegation are necessary for trustworthy, efficient model training. The spaCy+loss paradigm robustly separates learnable from non-learnable tokens, maximizing factual accuracy and minimizing factual entanglement, while retaining SLM NLU capabilities. The approach has practical scalability and theoretical significance in guiding SLM research towards fine-grained token-wise learning and intelligent deferral mechanisms.