Source-linked AI summary

Beyond Token-Level Guidance: Inference-Time Alignment of Specialized LLMs via Cross-Family Representation Steering

Jin Gan, Xin Li, Jun Luo

arXiv:2608.30319v1cs.CLcs.AIcs.LG

TL;DR

Specialization can degrade safety, while existing inference-time guidance may be unreliable because specialized and general guidance models have orthogonal competencies. CREST transfers safety directions through hidden-representation steering rather than token distributions, preserving generation boundaries and domain behavior. Across code, math, and medical domains, it improves weakened safety while preserving specialized capability and already-aligned safety, with reported gains up to 22.2%.

  • Problem

    Specialization fine-tuning can degrade safety, while existing inference-time methods may disrupt specialized generation because guidance models lack domain expertise.

  • Method

    CREST steers base-model hidden representations using safety directions extracted from a guidance model of any family, leaving token probabilities and EOS decisions unmodified.

  • Results

    CREST improves safety where specialization weakened alignment while preserving domain capability and safety in already well-aligned models, with gains up to 22.2%.

  • Takeaways & Limitations

    Representation-level steering addresses stop-token interference while retaining the specialized model’s control over generation boundaries.

  • Takeaways & Limitations

    CREST requires domain-appropriate paired safe/unsafe prompts, depends on detection-vector quality, and has so far been evaluated on three domains.

Abstract

from arXiv · show

Large language models (LLMs) finetuned for specialized domains represent crucial high-impact applications. Inference-time alignment improves safety degraded from specialization finetuning without requiring substantial computational resources, complementing finetuning-based methods with an easy-to-use, plug-and-play solution. However, existing inference-time methods fail to reliably improve safety without disrupting domain capability. We identify the root cause as complementary expertise orthogonality: specialized base models and general-domain guidance models have orthogonal competencies, making the guidance signal unreliable for specialized generation. This primarily manifests as stop token interference, where the guidance model's tendency toward continuation overrides the base model's decision to stop, burying correct answers under guidance-induced continuation. To address this problem, we propose CREST, an inference-time alignment method that steers base model hidden representations using safety directions extracted from a guidance model of any family, avoiding token-level structural limitations entirely. CREST improves safety where specialization has weakened it while preserving both domain-specific capability and the safety of already well-aligned models, outperforming baselines by up to 22.2\% on safety benchmarks. Our code is available at: https://github.com/DecayingSeart/CREST.

1 Introduction

Specialization gives LLMs domain expertise but can degrade safety, while existing inference-time guidance may disrupt specialized generation. CREST instead steers hidden representations using guidance-model safety directions, preserving stopping behavior and specialized capability.

  • Motivation: Specialized fine-tuning improves domain expertise but can substantially degrade safety alignment, creating risks for high-stakes deployments.Existing fine-tuning-based safety-preservation methods require substantial computational resources.
  • Motivation: Inference-time alignment offers a parameter-free alternative by using an aligned guidance model to steer an under-aligned specialized model during generation.This setup is practical when practitioners have one specialized model and rely on general-purpose aligned models.
  • Failure of Existing Methods: Existing token-level methods can suppress end-of-sequence signals, forcing continuation past correct domain answers and burying them under guidance-induced text.They apply the guidance distribution uniformly across token positions, including structural stop decisions.
  • Failure of Existing Methods: Complementary expertise orthogonality arises because specialized models retain domain expertise while general guidance models retain safety alignment, leaving their competencies orthogonal.The resulting failure is specific to specialized alignment and extends across inference-time methods regardless of guidance-signal quality.
  • CREST: CREST steers base-model hidden representations with safety directions from a guidance model of any family, leaving token probabilities—including EOS—unmodified.The base model retains control over stopping, and steering ends after the intervention window so domain generation resumes.
  • CREST: CREST improves safety where specialization weakened alignment while preserving specialized capabilities and safety in already well-aligned models, with gains up to 22.2%.The paper presents this as a plug-and-play complement to fine-tuning methods without family restrictions on the guidance model.

2 Related Works

Prior work aligns models during fine-tuning or generation, but inference-time token guidance can override specialized behavior. Single-model activation steering also assumes the target retains a reliable safety geometry, unlike the setting addressed here.

  • Inference-Time Alignment: Fine-tuning-based alignment can be undermined by adversarial examples and even benign downstream fine-tuning, motivating alternatives that avoid retraining.Inference-time alignment modifies outputs during generation without modifying model parameters.
  • Inference-Time Alignment: NUDGING replaces uncertain base-model tokens with guidance suggestions, while BLENDIN blends base and guidance distributions according to confidence.Both apply general-domain guidance at uncertain positions without distinguishing content tokens from structural tokens.
  • Limitations of Token Guidance: Token-level guidance can override specialized domain behavior because general-domain continuation preferences apply at positions where domain-specific stopping or content decisions are required.Uncertainty-based triggering can also fail to intercept harmful outputs generated confidently by specialized models.
  • Activation Steering: Activation-steering methods extract safety directions from the target model’s own activations, presupposing that it still encodes reliable safety geometry.The present setting instead motivates transferring safety information from a guidance model because specialization fine-tuning has degraded that geometry.

3 Problem Formalization

Specialized base and general guidance models have complementary expertise, making token-level alignment unreliable at domain-hard positions and especially vulnerable to stop-token interference.

  • 3.1 Complementary Expertise Orthogonality: Mb acquires domain expertise Ed through fine-tuning, while Mg retains safety alignment As but lacks exposure to the domain corpus.
  • 3.1 Complementary Expertise Orthogonality: Complementary expertise orthogonality combines expertise asymmetry, safety asymmetry, and non-dominance between specialized and guidance models.
  • 3.1 Complementary Expertise Orthogonality: At domain-hard positions, Mg’s signal can mislead because it lacks the domain knowledge needed to distinguish correct specialized responses from incorrect ones.
  • 3.1 Complementary Expertise Orthogonality: The characterization is falsifiable and descriptive rather than universal, with safety asymmetry expected to vary across domains and models.
  • 3.2 Structural Limitations of Token-Level Methods: Token-level methods apply guidance distributions uniformly, including at domain-hard content and structural positions where the base model decides whether to stop.
  • 3.2 Structural Limitations of Token-Level Methods: Stop-token interference occurs when general-domain continuation preferences override the specialized model’s EOS decision, burying correct answers under further generation.
  • 3.2 Structural Limitations of Token-Level Methods: 61.3% burial and 27.7% accuracy occur for Qwen3-8B despite 99.2% EOS suppression, while Llama-3.1-8B has 16.5% burial at 99.1% suppression.
  • 3.2 Structural Limitations of Token-Level Methods: Uncertainty-based triggering misses harmful completions when specialized models confidently generate domain-specific harmful content.

4 Method

CREST aligns safety representations across model families and detects threats from the specialized base model’s partial generation. It then applies adaptive, temporary hidden-state steering before returning control to unsteered domain-specific generation.

  • One-Time Setup: The one-time setup caches an alignment profile, selects guidance layers with the strongest safety separation, and maps them proportionally to base-model layers.In practice, k=1 usually suffices because safety signals concentrate in final aligned-model layers.
  • Representation Alignment: CREST transfers a guidance model’s safety representations into the base model’s hidden-state space using per-layer Procrustes alignment.The continuous projection remains applicable across different hidden dimensions, representation geometries, and tokenizer families.
  • Threat Detection: CREST detects threats using the specialized base model’s own representations from partial generations rather than relying on the guidance model’s domain knowledge.The detection vector is derived from safe and unsafe probe outputs, and the threshold is calibrated between their score distributions.
  • Per-Query Inference: For each query, an unsteered probe produces a threat score that determines adaptive steering strength through a threshold-centered sigmoid.Queries below the calibrated threshold receive no steering, preserving unmodified base-model generation on benign inputs.
  • Per-Query Inference: For detected threats, CREST computes a query-specific guidance safety direction, projects it into base-model space, and adds it to selected hidden states for Ti tokens.The safety direction connects a safe response prefix with the base model’s actual probe trajectory.
  • Per-Query Inference: After the intervention window, CREST removes its hooks and continues generation unsteered so the base model’s domain expertise governs the remaining completion.The algorithm explicitly resumes unsteered continuation after applying the temporary hidden-state correction.

5 Experiments

CREST is evaluated across code, math, and medical specializations using cross-family guidance, safety benchmarks, capability benchmarks, and sensitivity analyses. Results show safety gains where specialization weakened alignment, preservation on already aligned models, and negligible capability loss.

  • Experimental Setup: CREST is evaluated on code, math, and medical specialized models with distinct tokenizer families.The base models are Qwen2.5-Coder-7B-Instruct, Mathstral-7B-v0.1, and MedGemma-1.5-4B-it.
  • Experimental Setup: Llama-3.1-8B-Instruct provides cross-family guidance across all domains, while baselines use same-family guidance to reduce end-of-sequence suppression and capability disruption.This design tests transfer of safety representations across tokenizer-family boundaries while giving token-level baselines favorable conditions.
  • Evaluation: Safety is measured with insecure code rate, safety-sensitive math accuracy, and unsafe-medical refusal rate; capability is measured with HumanEval, GSM8K, and MedQA.The safety benchmarks contain 351 code, 450 math, and 466 medical examples; capability benchmarks use 114, 1,269, and 1,220 examples respectively.
  • Main Results: 22.2% relative reduction over the base model is achieved on code insecure rate, with CREST at 0.14 versus 0.18 for the base model.CREST also reaches 0.96 on math versus 0.92 for the base model and maintains 0.99 on medical versus 0.99 for the base model.
  • Sensitivity and Efficiency: Sensitivity analyses confirm robustness across hyperparameters and Llama, Qwen, and Gemma guidance families, while CREST incurs 1.3× base-only latency versus 2.8× for BlendIn and 7.0× for Nudging.CREST invokes the guidance model at most twice per query rather than once per token.

6 Conclusion

The conclusion identifies complementary expertise orthogonality and stop token interference as root causes of existing inference-time alignment failures. CREST addresses them through cross-family representation steering, improving safety while preserving capability and existing alignment.

  • Conclusion: Complementary expertise orthogonality causes existing inference-time alignment failures, primarily through stop token interference, across guidance-signal designs.Specialized models retain domain expertise while guidance models retain safety alignment, leaving their competencies orthogonal.
  • Conclusion: CREST avoids these failures by steering hidden representations rather than token distributions across model families.Experiments across code, math, and medical domains demonstrate safety improvements where specialization weakened alignment while preserving capability and already-aligned safety.

Limitations

CREST depends on domain-appropriate paired safe/unsafe setup prompts and is currently evaluated on only three domains, with indirect baseline capability comparisons and no demonstrated jailbreak robustness.

  • Scope and Assumptions: CREST requires a one-time setup phase with domain-appropriate paired safe/unsafe prompts.Performance depends on the quality of the detection vector, which may degrade when safety boundaries are ambiguous.
  • Scope and Assumptions: The method is evaluated on three domains, while broader specialized-domain and model-family coverage remains future work.Capability comparisons with token-level baselines are indirect because the methods use different generation backends.
  • Threat Model: CREST addresses safety degradation from benign specialization fine-tuning, whereas robustness to adversarial jailbreak attacks remains a distinct threat model.Jailbreak robustness is identified as a potential extension rather than an evaluated capability.

A.1 Training-Time Safety Preservation

Training-time safety-preservation methods modify alignment or fine-tuning, but they require training control, model-weight access, retraining resources, and tuning that may be unavailable to practitioners using third-party specialized models.

  • Training-Time Approaches: Training-time methods preserve safety through perturbation-aware alignment, memory replay, regularization, optimization changes, or data filtering.Examples include Vaccine, Booster, EWC, EMA momentum, LARF, and GradShield.
  • Limitations: These approaches require control over training, model-weight access, substantial retraining resources, and careful hyperparameter tuning.Such requirements can prevent practitioners with pretrained third-party specialized models or computational constraints from using them.

A.2 Capability-Safety Coupling

The guidance model lacks the specialized domain knowledge needed to distinguish domain-unsafe responses from domain-incorrect refusals. Consequently, interventions at domain-hard positions can unpredictably affect both capability and safety.

  • Because both capability and safety judgments depend on specialized domain knowledge, modifying domain-hard positions can disrupt both simultaneously.
  • At domain-hard positions, the guidance signal cannot distinguish correct specialized responses from factually incorrect refusals.
  • The guidance model’s intervention therefore has non-zero probability of degrading capability and safety together.

A.3 Sensitivity Analysis

CREST is generally stable across steering strength, probe length, steered layers, intervention length, and guidance-model choices, while threshold selection creates explicit safety-capability tradeoffs.

  • αmax produces stable results overall, with its optimal scale differing across domains because hidden-state scales vary.On code, αmax=1.6 achieves an insecure rate of 0.10 with Cap(∆)= −0.03.
  • Probe-length results are stable across domains; Tp=1 is selected for medical, while Tp=15 is selected for code and math.Longer probes provide additional context for threat assessment in code and math.
  • Safety remains stable across k=1, 2, and 3 steered layers, so k=1 is used by default to reduce computational overhead.
  • Moderate to long intervention lengths (Ti≥25) consistently rival or outperform very short interventions (Ti=5), particularly on code.
  • CREST maintains safety improvements across guidance-model configurations, with only modest performance variation.The reported robustness indicates that effectiveness is attributed to representation-space alignment rather than a particular guidance model.
  • Adaptive threshold calibration is a reliable default, whereas fixed thresholds can trade capability for safety and may require manual tuning.On code, τ=0.2 yields an insecure rate of 0.10 but HumanEval pass@1 of 0.05; on medical, fixed thresholds yield 0.99–1.00.

A.4 Runtime and Overhead

CREST adds less runtime overhead than token-level baselines and its judged medical results are sensitive to answer extraction. Full-response scoring leaves three regressions and three improvements, with the improvements reflecting intercepted harmful completions.

  • CREST adds 1.3× latency, or 6.0 versus 4.6 seconds per query, compared with 2.8× for BlendIn and 7.0× for Nudging.CREST performs at most two additional guidance-model forward passes per query, unlike token-level methods requiring repeated queries.
  • Trimmed extraction lowers judged medical safety from 0.99 to 0.96 and produces 24 apparent regressions, versus 3 under full-response scoring.Among the 24 trimmed-variant cases, 17 omit leading refusal clauses, 6 introduce invented follow-up exchanges, and 1 reflects a substantive response difference.
  • Math changes only marginally under trimmed extraction, from 0.96 to 0.95.
  • Under full-response scoring, CREST has 3 regressions and 3 improvements relative to the base model.The regressions are cases where the judge’s label could reasonably go either way, while the improvement cases involve harmful base-model completions intercepted by CREST.
Loading 2608.30319v1…