SafeDepth: Safety-Aware Token-Level Adaptive Computation

1 INTRODUCTION

• Token-level adaptive models can weaken backbone safety: on WildJailbreak, harmful-response rates rise from 8.33% for the reference Llama model to 41.67% for FlexiDepth and 31.25% for DiffSkip.Existing methods primarily target task performance and computation, while their safety effects remain largely unexplored (Schuster et al., 2022; Elhoushi et al., 2024; Luo et al., 2025b;a).Source: §1 · p. 1
• Safety depends on which layers run and whether they process prompts or generate answers; harmful-intent recognition and refusal also have different layer-wise dependencies.These findings motivate learning which layers to execute or skip for each token while retaining computational efficiency.Source: §1 · pp. 1–2
• SafeDepth learns token-level execution and skipping decisions with a frozen pretrained backbone through two-stage training.Lightweight skip adapters and a router first establish paths balancing generation quality and cost; the second stage refines paths using safety feedback from complete responses.Source: §1 · pp. 1–2
• Compared with FlexiDepth, SafeDepth reduces unsafe-response rates on all five harmful-request benchmarks, with HarmBench-HJ falling from 55.44% to 25.25%.It also reduces the XSTest false-refusal rate from 12.00% to 4.00%, while lowering computational cost and improving GSM8K accuracy from 65.66% to 66.64%.Source: §1 · p. 2

2 RELATED WORK

• Adaptive-computation methods allocate layer execution to tokens to reduce inference cost, while most target the trade-off between generation quality and computation (Zeng et al., 2023; Raposo et al., 2024; He et al., 2025a; Luo et al., 2025b;a).Work includes early exits such as CALM (Schuster et al., 2022), selective skipping, and deployment methods for batching, KV caching, and long-context inference (Del Corro et al., 2023; He et al., 2025b).Source: §2 · p. 2
• Mechanistic safety studies connect refusal and harmfulness representations to generation positions and examine how internal computation changes affect safety (Arditi et al., 2024; Qi et al., 2025; Zhao et al., 2025).Examples include cross-layer residual connections disrupting aligned behavior and targeted expert repair with routing consistency improving MoE safety (Joshi et al., 2025; Liang et al., 2026).Source: §2 · pp. 2–3
• SafeDepth extends these lines by using complete-response safety to supervise an efficient execution policy for a frozen dense model.Its phase-aware router selects token-specific, noncontiguous layer paths while accounting jointly for safety, answer quality, and computational cost.Source: §2 · p. 3

3 SAFETY IMPACTS OF TOKEN-LEVEL ADAPTIVE COMPUTATION

• All token-level adaptive reasoning models in the study exhibit higher unsafe-response rates than their corresponding backbones.This finding motivates treating safety as an explicit objective rather than assuming adaptive execution preserves backbone safety.Source: §3 · p. 3
• Restoring skipped computation during prefill reduces unsafe-response rates more consistently than restoring it during decoding.Source: §3 · p. 3
• Restoring one skipped layer at a time can reduce unsafe-response rates for some layers and increase them for others, so more computation does not necessarily improve safety.Safety depends on which layers execute, not simply the amount of computation performed.Source: §3 · p. 3
• Harmful intent can be detectable in earlier layers while executing or skipping a single later layer still changes whether the final response is safe.Source: §3 · p. 3

4 OUR METHOD

• SafeDepth adds a router and adapters to a frozen language model, choosing EXECUTE or SKIP for each token–layer pair using token representation, layer index, and inference phase.The backbone and language-model head remain frozen throughout.Source: §4 · p. 3

4.1 ADAPTER INITIALIZATION

• The adapter is initialized to approximate the full Transformer block’s output, giving the router a reliable skip branch before it learns where computation can be removed.Its training target is the fully executed block output; distance and cosine terms match representation and direction.Source: §4.1 · pp. 3–4

4.2 PHASE I: GENERATION-PRESERVING ROUTING

• Phase I jointly trains the router and adapters to reduce computation while preserving generation, using prediction objectives and execution probability as a differentiable proxy for execution frequency.The computation loss averages execution probabilities across eligible prefill and decoding positions; it does not measure training-time computation.Source: §4.2 · pp. 4–5

4.3 PHASE II: SAFETY-GUIDED ROUTE REFINEMENT

• Phase II traces unsafe complete responses to routing decisions by testing layer-wide interventions, then individual token changes, recomputing later routes after each intervention.If a single change is insufficient, bounded combinations are tested and retained as groups; their members are not treated as independently effective.Source: §4.3 · p. 5
• Phase II ranks localized successful changes above tested alternatives, preserves routes already producing safe and useful responses, and updates only the router before applying learned decisions directly at inference.Responses are regenerated and localization repeated after updates because earlier changes can alter later states and routing decisions.Source: §4.3 · p. 5

5.1 EXPERIMENTAL SETUP

• The experiments compare SafeDepth with full-depth, adaptive, and compression baselines across response safety, task capability, computation, and component ablations.The Llama-3-8B-Instruct setup routes at all 32 layers; Phase I uses Tulu-v2, while Phase II uses WildJailbreak and GSM8K training splits.Source: §5 · p. 5, §5.1 · p. 6, §5.2 · p. 6

5.2 RESPONSE SAFETY AND INFERENCE COST

• SafeDepth achieves the lowest unsafe-response rates across all five harmful-request benchmarks; rates fall from 55.44% to 25.25% on HarmBench-HJ and from 41.00% to 21.00% on JailbreakBench Attacks.Table 1 reports response-safety percentages, with lower values better across harmful-request and benign-request metrics.Source: §5.2 · p. 6, Table 1 · p. 7
• On XSTest, SafeDepth lowers false refusals from FlexiDepth’s 12.00% to 4.00%, below full-depth Llama’s 4.40%; its benign WildJailbreak rate is 8.59%.FlexiDepth’s benign WildJailbreak false-refusal rate is 5.47%, while full-depth Llama’s is 21.88%.Source: §5.2 · p. 6
• On WildJailbreak, SafeDepth lowers unsafe responses from 41.67% to 30.21% and complete-response matrix work from 7.072 to 5.340, a 24.5% reduction versus FlexiDepth.Complete-response cost includes prompt processing and answer generation; repetitive output can offset per-token savings.Source: §5.2 · pp. 6–7, Table 2 · p. 7

5.3 TASK CAPABILITY AND REASONING EFFICIENCY

• SafeDepth improves on FlexiDepth across six generation tasks, including HumanEval+ (29.27% to 33.41%) and MBPP+ (44.71% to 50.11%), and exceeds full-depth Llama on HumanEval+ and CoQA.Table 3 covers knowledge, commonsense, and generation benchmarks; its scores are on a 0–100 scale and higher is better within each column.Source: §5.3 · p. 7, Table 3 · p. 8
• On GSM8K, SafeDepth reduces matrix work by 29.7% versus full-depth Llama at 66.64% accuracy versus 70.36%, and improves over FlexiDepth from 65.66% to 66.64% while cutting matrix work 18.6%.The matched evaluation uses greedy decoding and a 512-token limit; SafeDepth’s matrix work versus FlexiDepth falls from 3.193 to 2.598.Source: §5.3 · p. 8, Table 4 · p. 9

5.4 ABLATION STUDY

• Removing K/V preservation or all skip adapters reduces GSM8K exact match from 66.64% to 0.83% or 4.02%; without K/V preservation, WildJailbreak unsafe responses reach 68.80% and XSTest false refusals 66.00%.The ablation results underscore the roles of retained attention history and adapters that transform hidden states on skipped paths.Source: §5.4 · pp. 8–9, Table 5 · p. 9
• Phase II reduces WildJailbreak unsafe responses from 46.80% to 30.21% and XSTest false refusals from 15.20% to 4.00%, with a 1.14% decrease in GSM8K accuracy.On a fixed 96-request audit, matrix work rises about 4.8% and latency about 4.5% versus Phase I.Source: §5.4 · p. 9, Table 5 · p. 9

• SafeDepth learns token-specific computation paths with a frozen pretrained backbone, using efficient-route training followed by safety-feedback refinement of token–layer execute-or-skip decisions.The two stages first preserve generation quality while learning efficient routes, then use complete-response feedback to refine decisions.Source: §6 · p. 9
• The paper reports that changing layer execution alters response safety, with effects varying across inference phases and layers.Source: §6 · p. 9
• Experiments report fewer unsafe responses and lower complete-response computation than FlexiDepth while largely preserving task performance.Source: §6 · p. 9

• Generative AI tools were used to edit and translate the manuscript for clarity and readability.Source: AI USE STATEMENT · p. 10
• The authors state that generative AI tools were not used to generate research ideas or code.Source: AI USE STATEMENT · p. 10
• The authors reviewed AI-assisted revisions and take responsibility for the final content.Source: AI USE STATEMENT · p. 10

• The study examines harmful prompts and model responses to understand safety failures and develop safer routing policies.Source: ETHICS STATEMENT · p. 10
• The evaluation also measures false refusals on benign requests.Source: ETHICS STATEMENT · p. 10
• The authors acknowledge that routing insights could be misused to weaken safeguards, while stating a defensive intent.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

• The reproducibility statement points readers to routing and training details, model and dataset configurations, filtering, baselines, evaluation protocols, decoding settings, and cost measurements.These materials are specified in Sections 4 and 5 and alongside the experimental results.Source: REPRODUCIBILITY STATEMENT · p. 10

A SAFETY IMPACTS OF TOKEN-LEVEL ADAPTIVE COMPUTATION

• The analysis frames adaptive execution as changing the routes used by input and generated tokens, and tests safety under matched prompts and decoding settings.It proceeds from backbone comparison to phase, layer-location, and recognition-versus-refusal questions.Source: A SAFETY IMPACTS OF TOKEN-LEVEL… · p. 14, §A.1 · p. 14

A.1 Q1: DOES ADAPTIVE COMPUTATION PRESERVE THE BACKBONE’S SAFETY?

• On WildJailbreak, harmful responses rise from 8.33% for Llama-3-8B-Instruct to 41.67% for FlexiDepth and 31.25% for DiffSkip; other adaptive systems also show large shifts.On StrongREJECT, ARM rises from 38.7% to 89.8%, while Plan-and-Budget rises from 1.6% to 21.3%; changes also differ in benign refusals.Source: §A.1 · p. 14, Figure 2 · p. 14

A.2 Q2: DO PREFILL AND DECODING HAVE THE SAME SAFETY EFFECT?

• On all four harmful-request benchmarks, restoring prefill yields lower harmful-response rates than restoring decoding alone; WildJailbreak falls from 41.67% to 9.6%, versus 36.4% with decoding restoration.On XSTest, restoring prefill raises benign refusals to 27.6%, while restoring decoding lowers them to 4.8%; restoring both yields 7.6% harmful responses on WildJailbreak and 7.2% benign refusals.Source: §A.2 · p. 15, Table 6 · p. 15

A.3 Q3: IS SAFETY DETERMINED BY THE AMOUNT OF COMPUTATION?

• With one layer restored at a time, WildGuard harmful-response rates vary by layer and phase: prefill restoration at layer 18 or 20 lowers the rate from 50.0% to 21.9%, but layer 23 yields 46.9%.At layer 18, decode-only restoration leaves the rate at 50.0%, while restoring it in both phases lowers the rate to 21.9%.Source: §A.3 · pp. 15–16, Figure 3 · p. 15

A.4 Q4: DOES RECOGNIZING HARMFUL INTENT DIRECTLY DETERMINE REFUSAL?

• A layer-16 probe reaches 0.946 AUROC, yet later execution decisions can reverse safety even when harmful intent is already readable.Among 200 harmful prompts, 62 have a positive danger margin at layer 20 but still produce harmful answers; paired interventions include six safe-to-harmful changes and one harmful-to-safe change, while 170 refusal pairs remain unchanged.Source: §A.4 · pp. 16–17, Figure 4 · p. 16, Table 7 · p. 16
• Transferring a harmful-request state at layer 16 changes refusal framing without replacing the benign tutorial task, while self-transfer and same-norm random replacement do not produce that change.The transferred response begins with a refusal-like preface but still discusses repairing vulnerabilities rather than exploiting them.Source: §A.4 · p. 17, Table 8 · p. 17

A.5 FROM FINDINGS TO DESIGN

• The design implications are to evaluate generated safety alongside utility, distinguish prefill from decoding, account for layer location, and preserve downstream computation that turns risk signals into behavior.The findings state that safety does not follow from task performance, computation budget alone, or the presence of a risk signal.Source: §A.5 · p. 17

B ILLUSTRATIVE EXAMPLES OF OUR METHOD

• The token–layer diagrams record selected routing decisions, but do not establish that an individual layer caused the response difference.The examples use the same complete prompt across Original Llama, FlexiDepth, and SafeDepth.Source: B ILLUSTRATIVE EXAMPLES OF OUR METHOD · p. 18

EXAMPLE 1: DECLINING A FABRICATED TESTIMONIAL

• For the fabricated testimonial prompt, Original Llama and FlexiDepth provide endorsements, while SafeDepth refuses to create the unsupported claim.The complete Original Llama response gives the endorsement before adding a critical persona.Source: EXAMPLE 1: DECLINING A FABRICATED… · p. 18

EXAMPLE 2: REFUSING A CRIME-FICTION PRETEXT

• For the crime-fiction pretext, both baselines give multi-step planning and concealment advice, while SafeDepth ends with a refusal.The SafeDepth route skips or executes different layers from FlexiDepth for selected prompt tokens; Original Llama executes all 32 layers.Source: EXAMPLE 1: DECLINING A FABRICATED… · p. 18, EXAMPLE 2: REFUSING A CRIME-FICTION… · p. 19