PaperlayerView exact revision on arXiv

Source-linked AI summary

ALRA: Adaptive Local Relational Alignment for Logit-Based Pre-training Distillation of Autoregressive Language Models

Quang Hoang Trung, Quang Huu Hieu, Nguyen Van Hoang Phuc, Vo Nguyen Le Duy

arXiv:2609.03355v1stat.MLcs.CLcs.LG

TL;DR

Global next-token distillation overlooks preferences among likely alternatives, while existing local methods select candidates from only the teacher or student. ALRA combines student-proposed tokens with teacher guidance and aligns both local distributions and token-pair preferences. Across nine zero-shot benchmarks, it outperforms the strongest compared distillation baseline for both student sizes.

  • Problem

    Full-vocabulary distillation overlooks relative preferences among likely alternatives, while teacher-only or student-only local selection can omit useful candidates or rely on inaccurate rankings.

    Source: §1 · pp. 1–2

  • Method

    ALRA combines student-proposed candidates with a teacher top-token anchor, adapts the local-set size, and aligns local/rest distributions and relative preferences among token pairs.

    Source: §1 · pp. 2–3, §3.3 · p. 9

  • Results

    36.62% and 37.40% average accuracy across nine zero-shot benchmarks, exceeding the strongest compared distillation baseline by 0.94 and 0.83 percentage points, respectively.

    Source: §1 · p. 3

  • Takeaways & Limitations

    ALRA ranks first or second on six of nine benchmarks at 200M parameters and seven at 500M, with gains over pre-training without KD on eight and nine benchmarks, respectively.

    Source: §4.2 · p. 14

    ALD is generally not algebraically identical to full-vocabulary forward KL because its conditional terms use unit coefficients.

    Source: §3.3 · p. 9

Abstract

from arXiv · show

Logit-based knowledge distillation for autoregressive language models usually aligns teacher and student next-token distributions over the entire vocabulary. However, this global objective overlooks relative preferences among likely token alternatives. Existing local approaches often select candidate tokens from either the teacher or the student alone. Teacher-only selection can miss tokens that the student considers likely, while student-only selection can rely on an inaccurate ranking early in training. We propose Adaptive Local Relational Alignment (ALRA), a position-specific framework combining student proposals with teacher guidance. At each valid prediction position, the student proposes likely tokens, while the teacher's most probable token is included as an anchor. ALRA adjusts the number of selected tokens according to how broadly the teacher distributes probability within this candidate set relative to the current batch. Adaptive Local Divergence retains the mass-matching term and separately matches the relative token distributions within the selected and remaining vocabulary regions. Unlike the exact full-vocabulary decomposition, it replaces the teacher-mass coefficients of the two conditional terms with unit coefficients, preventing either term from being downweighted solely because its region has low teacher probability. Student-Weighted Pairwise Relational Alignment emphasizes high-probability token pairs with small student probability gaps and gives less weight to unlikely or clearly separated pairs. Experiments on The Pile with randomly initialized 200M- and 500M-parameter students across nine zero-shot benchmarks yield average accuracies of 36.62% and 37.40%. ALRA exceeds the strongest competing distillation baseline by 0.94 and 0.83 percentage points and improves over pre-training without distillation by 2.31 and 2.91 points, respectively.

1. Introduction

Full-vocabulary distillation treats every token within one global objective, while local approaches can rely too heavily on either teacher or student rankings. ALRA combines student-proposed candidates with teacher guidance and pairs adaptive regional matching with relational supervision.

1. Introduction

  • Full-vocabulary forward KL uses one global objective despite context-varying teacher uncertainty, while teacher-only and student-only selection each have distinct shortcomings.Teacher-only selection may miss student-likely tokens; student-only selection can be unreliable early in training, and fixed local budgets do not reflect context-dependent alternatives.Source: §1 · pp. 1–2
  • ALRA is a position-specific logit-distillation framework combining adaptive token selection, region-aware divergence, and student-dependent pairwise relational supervision.The method updates only the student model.Source: §1 · p. 2
  • At each prediction position, ALRA uses student-proposed candidates plus the teacher top-1 token, then adapts the local-set size to teacher support relative to the batch.The local set is selected from the teacher’s highest-probability tokens within the anchored candidate set.Source: §1 · pp. 2–3
  • Adaptive Local Divergence matches local and rest regions alongside their probability mass, while SWPRA emphasizes plausible token pairs that the student currently separates weakly.ALD gives unit coefficients to its two conditional divergences so neither is scaled down solely by its region’s teacher probability mass.Source: §1 · pp. 2–3
  • 36.62% (200M) and 37.40% (500M) average accuracy across nine zero-shot benchmarks exceed the strongest competing distillation baseline by 0.94 and 0.83 percentage points, respectively.Students were randomly initialized and distilled from a frozen Qwen1.5-1.8B teacher on The Pile.Source: §1 · p. 3

2. Related Work

Related work spans output-logit, representation-level, adaptive, tail-aware, sequence-level, data-centric, and policy-based approaches to transferring teacher knowledge. ALRA is positioned at the intersection of adaptive local selection, local-and-rest supervision, and pairwise relations among vocabulary entries.

2. Related Work

  • White-box logits-based KD matches full next-token distributions, whereas hint-based KD also supervises intermediate representations; this paper uses output-level signals because teacher and student hidden sizes differ.The students share the teacher’s output vocabulary, avoiding explicit alignment of intermediate representation spaces.Source: §2 · p. 3
  • Prior relational KD preserves example relations in representation space, while output-level methods model class dependencies.ALRA follows the output-level direction by modeling relations among vocabulary entries.Source: §2 · p. 4
  • Adaptive and selective methods allocate supervision across samples, positions, or token difficulty, while BiLD filters long-tail logits and TAD increases the contribution of lower-probability predictions.These approaches motivate controlling where and how teacher information is transferred across vocabulary regions.Source: §2 · p. 4
  • Sequence-level, data-centric, and policy-based approaches alter generated sequences, pre-training data, or student trajectories; the paper treats them as complementary rather than direct baselines.The authors note these strategies could be combined with their token-level logit objective in future work.Source: §2 · pp. 4–5
  • ALRA targets the stated intersection of position-specific adaptive local budgets, supervision of local and complementary vocabulary regions, and pairwise relational weighting.The paper compares complete objectives under a shared pre-training pipeline and reserves mechanism-specific conclusions for controlled studies.Source: §2 · p. 5

3. Proposed Approach

The proposed approach aligns teacher and student at each valid prediction position through an adaptive, teacher-anchored local token set, mass and conditional divergence terms, and student-weighted pairwise alignment. It computes losses using a batch-adaptive local budget and can include a causal language-modeling term when labels are available.

3.2. Teacher-Anchored Adaptive Local Token Set

  • ALRA selects each local set from student-proposed tokens plus the teacher top-1 anchor, sizing it by teacher effective support relative to the batch and ranking candidates by teacher probability.The selected tokens are drawn from the anchored candidate set; unselected vocabulary tokens remain in the rest region.Source: §3 · p. 5, §3.2 · pp. 6–8

3.3. Adaptive Local Divergence

  • Adaptive Local Divergence retains local-versus-rest mass matching but gives both conditional divergences unit coefficients, avoiding downweighting low-teacher-mass regions.ALD is a distinct region-aware objective, not an algebraically exact rewriting of full-vocabulary forward KL.Source: §3.3 · p. 9

3.4. Student-Weighted Pairwise Relational Alignment

  • Student-Weighted Pairwise Relational Alignment weights pairs within each position-specific local set by student probability mass and gap, favoring likely alternatives with similar student probabilities.The pairwise KL targets remain teacher-to-student; the student-derived weights change only each pair’s contribution.Source: §3.4 · pp. 9–10

3.5. Training Objective

  • ALRA retains the mass term and gives its local-conditional and rest-conditional components unit coefficients; λ_pair controls the pairwise relational term.The comparison variants replace SWPRA with PRA or PRA-U, without introducing additional region-level weighting hyperparameters.Source: §3.5 · p. 11
  • When ground-truth next-token labels are available, the objective can include a causal language-modeling term weighted by λ_CE ≥ 0.The student’s temperature-one next-token distribution is defined for this setting.Source: §3.5 · p. 11
  • Algorithm 1 averages candidate-conditioned teacher effective support across the batch to set the local budget, then averages position-level ALD and SWPRA objectives.It forms a teacher-anchored candidate set and local–rest partition at each valid prediction position; the causal language-modeling term is added afterward when used.Source: §3.5 · p. 11

4. Experiments

The experiments evaluate ALRA in matched from-scratch pre-training and analyze choices for its local token budgets, anchoring, and pairwise weighting. They report average zero-shot and held-out accuracy across two student sizes to compare these design variants.

4.1. Experimental Setup

  • With the approximately 1B-token budget, randomly initialized 200M and 500M Qwen students trained from scratch are evaluated zero-shot on nine held-out benchmarks against a fixed Qwen1.5-1.8B teacher.Methods share data order, token budget, sequence length, and optimization schedule; reported checkpoints are at the end of the budget, not selected by downstream scores.Source: §4.1 · pp. 11–13

4.2. Main Results on Held-Out Benchmarks

  • 36.62% for the 200M student exceeds TAD by 0.94 points, while 37.40% for the 500M student exceeds ATKD by 0.83 points, leading both nine-benchmark averages.ALRA improves over pre-training without KD by 2.31 and 2.91 points for the 200M and 500M students, respectively; it ranks first or second on six and seven benchmarks, respectively.Source: Table 1 · p. 13, §4.2 · pp. 13–14

4.3. Analysis of Adaptive Local Budget Allocation

  • Adaptive ALRA scores 39.89 and 41.09 validation accuracy, beating the best fixed-budget results by 0.50 and 0.36 points for the 200M and 500M students.The best fixed budgets differ by size: d=15 yields 39.39 for 200M, while d=19 yields 40.73 for 500M. Among non-maximum budgets, the late-training distribution peaks around d_u≈9–10; the share assigned d_max rises from 48.1% to 51.7%.Source: §4.3 · pp. 14–16, Table 2 · p. 15, Figure 3 · p. 16

4.4. Effect of Teacher Top-1 and Ground-Truth Anchoring

  • Teacher top-1 anchoring improves validation averages over unanchored ALRA-N by 1.03 points for 200M and 1.02 for 500M, whereas adding both anchors underperforms ALRA by 0.68 and 0.57 points.Ground-truth anchoring changes the 200M average from 38.86 to 38.87 and the 500M average from 40.07 to 40.74. The teacher top-1 and ground-truth tokens coincide at 53% of sampled positions; when they differ, including both can displace student proposals, and low-teacher-probability ground-truth tokens may not remain in the final local set.Source: Table 3 · p. 17, §4.4 · p. 17, Figure 4 · p. 18
  • Teacher top-1 anchoring gave the highest validation average for both student sizes among four variants; ground-truth anchoring, alone or combined with it, performed worse.The final ALRA configuration therefore uses only the teacher top-1 token as the candidate-set anchor.Source: §4.4 · p. 18

4.5. Effect of Pairwise Weighting

  • 34.91% to 36.56% and 35.77% to 37.09%: PRA-U improved the average for 200M and 500M by normalizing pair weights, making the pairwise term an average rather than a sum.Normalization gives a more comparable pairwise signal across positions with different numbers of local pairs.Source: §4.5 · p. 18, Table 4 · p. 18
  • SWPRA gives greater normalized weight to plausible token pairs with more student probability mass and smaller probability gaps, while downweighting unlikely pairs or pairs already clearly separated by the student.This emphasizes local comparisons where the student's relative preference is still weak.Source: §4.5 · pp. 18–19
  • 36.62% and 37.40%: SWPRA raised the held-out average from 36.56% to 36.62% for 200M and from 37.09% to 37.40% for 500M, and was used in final ALRA.The compared variants share the same local token pairs and pairwise KL; Table 4 isolates the pair-weighting choice.Source: Table 4 · p. 18, §4.5 · p. 19

5. Conclusion

ALRA combines student-proposed tokens with teacher guidance in a position-specific local set, then applies region-aware divergence and pairwise relational supervision. Across nine zero-shot benchmarks, it outperforms the strongest competing distillation baseline for both student sizes, within a white-box, shared-vocabulary setting.

5. Conclusion

  • 36.62% and 37.40% average accuracies for randomly initialized 200M- and 500M-parameter students exceed the strongest competing distillation baseline by 0.94 and 0.83 percentage points, respectively.The students were distilled from a frozen Qwen1.5-1.8B teacher and evaluated across nine zero-shot benchmarks.Source: §5 · p. 19
  • ALRA combines student-proposed tokens with teacher guidance to construct a position-specific local token set.Adaptive Local Divergence aligns probability mass and conditional distributions in the local and remaining vocabulary regions; Student-Weighted Pairwise Relational Alignment emphasizes local pairs with high student probability mass and small probability gaps.Source: §5 · p. 19
  • The formulation is limited to white-box logit distillation with a shared teacher–student output vocabulary.Different tokenizers or more limited teacher-output access would require another way to define and align the local token space.Source: §5 · p. 19
  • The work focuses on token-level pre-training distillation for autoregressive language models; combining ALRA with sequence-level distillation or other training stages remains open.Source: §5 · p. 19

A. Derivation of the Local–Rest KL Decomposition

The appendix relates the local–rest partition to the full-vocabulary forward KL and explains how ALD modifies that decomposition. It also describes the teacher and randomly initialized students used in the experiments, whose architectures create different capacity gaps.

A. Derivation of the Local–Rest KL Decomposition

B. Model Configurations

  • The teacher and students share a Qwen-family architecture, tokenizer, and 151,936-token vocabulary; the students are randomly initialized, with the 200M model reducing depth and width and the 500M model reducing width.The teacher remains fixed, and the released student weights are not used.Source: B. Model Configurations · pp. 21–22, Table B.1 · p. 22

C. Training Data Construction

Training text comes from a filtered version of The Pile, with the first two million streamed records retained. Shared-tokenizer encoding, document-boundary markers, and boundary-aware chunking produce aligned next-token examples up to the training input length.

C.1. Source Corpus and Document Selection

  • The corpus is Pile Uncopyrighted, a filtered version of The Pile, and the first 2,000,000 records in stream order are retained.The selected records contain approximately 2.8B tokenizer tokens before boundary EOS insertion and final budget truncation; only each record’s text field is used.Source: §C.1 · p. 22

C.2. Tokenization and Boundary-Aware Chunk Construction

  • Text is encoded with the shared Qwen tokenizer, marked with an EOS token, and chunked into examples of up to 513 token IDs that provide 512 inputs and 512 next-token targets.When a chunk reaches maximum length, boundary-aware construction searches backward for an acceptable EOS, newline, or sentence-ending boundary and returns later tokens to the next example’s buffer.Source: §C.2 · p. 22

C.3. Data Ordering and Token Budgets

The shorter controlled runs use prefixes of the same precomputed training stream, while compared methods follow the same example order. Nominal token budgets are derived from updates and batch dimensions, and the evaluation suite uses designated held-out splits.

C.3. Data Ordering and Token Budgets

  • The approximately 550M-token controlled run is an exact prefix of the approximately 1B-token stream, and methods at the same budget see the same precomputed sequence order without shuffling.Source: §C.3 · p. 22
  • 67,200 updates equal 550,502,400 nominal model-input tokens, while 126,400 equal 1,035,468,800; ignored padded targets can reduce positions contributing to the objective.Source: §C.3 · p. 22
  • Held-out evaluation uses official test splits where usable and public validation splits otherwise; reserved validation examples are used only for the controlled 550M-token analyses.Source: Table D.1 · p. 23

D.1. Benchmark Tasks, Splits, and Metrics

The section defines the held-out and validation splits used across nine benchmarks, including a fixed fallback when official test labels are unusable. Figure D.1 distinguishes validation sets for controlled studies from held-out sets for the benchmark comparison.

D.1. Benchmark Tasks, Splits, and Metrics

  • Official test data serve as held-out sets when usable; otherwise, public validation splits are used for HellaSwag, WinoGrande, PIQA, and SocialIQA.LAMBADA-OpenAI, OpenBookQA, ARC-Easy, ARC-Challenge, and StoryCloze-2016 use their official test splits.Source: §D.1 · p. 23
  • Separate validation sets are reserved for controlled analyses only on HellaSwag, ARC-Easy, and PIQA; the other six tasks have none.ARC-Easy uses its official validation split, while HellaSwag and PIQA use labeled training examples after their public validation splits become held-out sets.Source: §D.1 · p. 23
  • The figure is descriptive, not a weighting rule: the “Avg.” score weights each benchmark equally, so larger HellaSwag and PIQA sets do not dominate.Table D.1 reports the task identifiers and example counts used in the protocol.Source: §D.1 · p. 23
  • Figure D.1 uses a logarithmic vertical axis, with blue bars for controlled-study validation sets and orange bars for nine-benchmark held-out sets.Blue bars appear only for HellaSwag, ARC-Easy, and PIQA.Source: §D.1 · p. 23, Figure D.1 · p. 24

D.2. Controlled Mechanism Analysis

Controlled mechanism analyses compare methods on the same three validation sets after the same shorter training budget. The zero-shot benchmark suite scores tasks without labeled downstream updates and gives each benchmark equal weight in the final average.

D.2. Controlled Mechanism Analysis

  • Tables 2 and 3 compare the fixed-budget sweep, anchoring variants, and full ALRA on the same ARC-Easy, HellaSwag, and PIQA validation sets after about 550M tokens.These mechanism-oriented analyses are separate from the nine-task held-out results in Table 1; Table 4 reports the nine-task held-out average after about 1B tokens.Source: §D.2 · p. 24

D.3. Zero-Shot Scoring and Averaging

  • Zero-shot scoring uses default LM Evaluation Harness configurations without labeled downstream examples updating student parameters.The eight selection tasks use multiple-choice accuracy; LAMBADA-OpenAI uses exact next-word accuracy.Source: §D.3 · p. 24
  • The final average assigns each of the nine benchmarks weight 1/9, regardless of its example count.The evaluation uses the latest available LM Evaluation Harness commit, with its identifier and exported task configurations released in the evaluation code.Source: §D.3 · p. 24

E.1. Optimization Configuration

The experiments use a shared optimization setup and report end-of-budget checkpoints rather than selecting by downstream performance. ALRA’s local-set and pairwise settings are specified separately, including which hyperparameter controls probability-gap weighting.

E.1. Optimization Configuration

  • Students use AdamW with a 16-sequence effective batch, gradient clipping at 0.5, and a cosine learning-rate schedule after 512 warmup updates.The learning rate warms to 6 × 10^-4 and ends at 6 × 10^-5; controlled runs stop at the corresponding intermediate point of the shared 126.4K-update schedule.Source: §E.1 · p. 24
  • Reported results use end-of-budget checkpoints—67.2K updates for controlled analyses and 126.4K for approximately 1B-token experiments—not checkpoints selected by downstream scores.The runs include no across-seed statistics, so the authors do not claim statistical significance for small numerical differences.Source: §E.1 · p. 25

E.2. ALRA Hyperparameters

  • ALRA uses d_min = 3 and d_max = 25, with fixed-budget comparisons at d ∈ {3, 7, 9, 15, 19, 25}.Distillation and pairwise temperatures are 1.0, λ_pair = 1.0, ε = 10^-6, and the reported comparison sets λ_CE = 0.Source: §E.2 · p. 25
  • The adaptive budget uses teacher probabilities conditioned on the student-proposed, teacher-anchored set, normalized by batch-average effective support—not student full-vocabulary entropy.γ = 5 controls the SWPRA probability-gap factor, not the complete pairwise-loss coefficient λ_pair.Source: §E.2 · pp. 25–26

E.3. Baseline Implementations and Adaptations

The appendix describes baseline adaptations and the added runtime and memory costs of ALRA. Comparisons use shared architectures, initialization, data, and optimization settings, while methods differ in their objectives and method-specific settings.

E.3. Baseline Implementations and Adaptations

  • Baseline students share ALRA’s architectures, initial parameters, processed corpus, data order, sequence length, update budget, and optimization configuration.Methods differ in objective construction and method-specific settings; reported loss weights specify the distillation-to-language-modeling mixture.Source: §E.3 · p. 26, Table E.3 · p. 27
  • Vanilla KD matches full teacher and student vocabulary distributions with forward KL, while PD uses temperature-scaled, truncated teacher logits and a fixed pure-distillation weight.Vanilla KD uses equal KL and causal-LM weights at temperature 0.5; PD’s reported configuration is an adaptation, not the full WSD-α schedule.Source: §E.3 · p. 27
  • ATKD splits tokens into hard and easy halves, applying both target- and diversity-oriented KD to the hard half and omitting the target term for the easy half.The reported adaptation uses equal subset weights and does not adapt the number of local vocabulary alternatives.Source: §E.3 · p. 27
  • RLD refines teacher logits using ground-truth labels at each valid prediction position; LDRLD transfers relations over top-d student logits with d = 15 for the large vocabulary.RLD is not attributed cross-layer supervision or an adaptive local budget; LDRLD adapts an image-classification method to vocabulary indices and retains a 0.5 language-modeling weight.Source: §E.3 · pp. 27–28
  • TAD matches a teacher-preferred top-10 head and normalized tail, but this adaptation omits its hidden-state cosine loss and teacher-derived attention initialization.The logit-level terms are combined with causal-LM training at equal weights because teacher and student hidden sizes differ.Source: §E.3 · p. 28
  • BiLD aligns teacher-led and student-led pairwise logit differences, applied directly to randomly initialized students rather than its original three-stage fine-tuning pipeline.The reported adaptation uses k = 8 and temperature 3.0, so it is not a reproduction of BiLD’s original setting.Source: §E.3 · p. 28

E.4. Hardware, Runtime, and Pairwise Cost

  • SWPRA adds O(d_u^2) pairwise construction within the adaptive local set, bounding the added pairwise work rather than scaling it with vocabulary size.At d_max = 25, the maximum pair count is 300; the reported overhead measurements use a single RTX 4090 with batch size 16 and maximum sequence length 512.Source: §E.4 · p. 28
  • For the 500M student, ALRA takes 2.75 seconds per step versus 2.23 seconds with Vanilla KD, about 23% more, and peak memory rises from 13.0 GB to 16.5 GB.For the 200M student, peak memory rises from 12.0 GB to 14.0 GB; on one RTX 4090, full-run wall-clock time is about three days for 200M and four for 500M.Source: §E.4 · p. 28

Related papers

Loading 2609.03355v1…