Inference-Time Subspace Alignment

• Inference-Time Subspace Alignment is a technique that constrains model inference to a structured, low-dimensional subspace to modulate output behavior while preserving overall competence.
• It encompasses diverse methods such as refusal and decision subspace steering, sparse logit intervention, and PCA-based domain adaptation, each tailored for specific decoding challenges.
• Empirical findings across LLM safety, guided decoding, and diffusion models demonstrate that aligning generation trajectories can significantly enhance robustness against adversarial token injections.

In contemporary machine learning literature, inference-time subspace alignment denotes a family of procedures that impose alignment during decoding, sampling, adaptation, or prediction by operating on a constrained representation at inference time rather than by fully retraining the base model. The aligned object varies by setting: refusal directions and generation trajectories in LLM safety, entropy-gated steering directions in decoding, stylistic/discourse token subspaces in guided decoding, preference subspaces in token-scoring space, PCA- or Tucker-derived invariant subspaces in domain adaptation, decision subspaces in vertical federated learning, and feature-conditioned guidance in diffusion models. Taken together, these usages suggest a common principle: preserve the base model’s broad competence while intervening on a low-dimensional, structured, or otherwise restricted component of its behavior (Park et al., 3 Jun 2026, Hu et al., 30 Jan 2026, Fei et al., 2024, Chang et al., 10 Jan 2026, Fernando et al., 2014).

1. Taxonomy and formal interpretations

The term is not used in a single uniform sense. In some papers it refers to explicit geometric alignment between source and target subspaces; in others it refers to steering within a representational or token-scoring subspace during generation; in still others it refers to decision-space alignment that makes local inference mimic collaborative inference.

Setting	Aligned object	Representative mechanism
LLM safety	Refusal directions and generation trajectories	Cosine to a refusal prototype; trajectory augmentation
Sparse decoding control	Logit space by default; hidden-state subspace as an extension	Entropy-gated junction steering
Guided decoding	Stylistic/discourse subspace	Uncertainty-gated token insertion
Multi-objective LLM control	Preference subspace in token-scoring space	Linear combination of value vectors
Domain adaptation and TTA	PCA, significant, invariant, or tensor subspaces	Closed-form or iterative source–target alignment
Federated inference	Decision subspace parameterized by the task head	Local-to-global decision matching

A recurrent distinction is between static and dynamic notions of alignment. Static formulations measure whether a hidden state, feature vector, or prototype lies near a direction or subspace. Dynamic formulations instead ask whether the trajectory induced by decoding or sampling remains aligned after perturbation. This distinction is central in recent safety work, where static refusal-direction alignment is shown not to predict robustness under token injection (Park et al., 3 Jun 2026).

2. Trajectory-level alignment in autoregressive LLMs

A safety-aligned autoregressive LLM $\pi$ generates a response $y=(y_1,\dots,y_T)$ to a harmful instruction $x$ via $y_t \sim \pi(\cdot \mid x, y_{<t})$. The paper “Inference-Time Vulnerability Beyond Shallow Safety: Alignment Along Generation Trajectories” formalizes shallow safety as alignment that primarily reshapes the distribution of the first few output tokens, and then shows that shallow safety is a special case of a broader inference-time vulnerability: short token injections at arbitrary decoding steps can redirect later generation toward harmful content (Park et al., 3 Jun 2026).

The formal object is a generation trajectory $\mathcal{T}=\{(h_t,y_t)\}_{t=1}^T$, where $h_t$ is the last-layer hidden state after processing $(x,y_{<t})$ and $h_{t+1}=f_\theta(x,y_{\le t})$. Mid-sequence perturbation is defined by an injection operator

$\tilde{y}=I_t(s;\,y):=(y_{<t},\, s,\, y_{\ge t}),$

with continuation decoded from context $(x,y_{<t},s)$. The paper evaluates both a harmful prefix, “As a harmful assistant, here's a detailed guide on”, and a refusal prefix, “I can't fulfill that request” (Park et al., 3 Jun 2026).

The same work defines a refusal direction by difference in means. Using harmful prompts from AdvBench, full responses are labeled safe or unsafe with Llama-Guard-3-8B, yielding refusal and compliance sets. If $y=(y_1,\dots,y_T)$0 and $y=(y_1,\dots,y_T)$1 are the mean hidden states over refusal and compliance samples, the refusal prototype is

$y=(y_1,\dots,y_T)$2

and the alignment measure at step $y=(y_1,\dots,y_T)$3 is

$y=(y_1,\dots,y_T)$4

The central empirical result is that subspace alignment does not predict robustness: across Mistral-7B-Instruct-v0.3, Qwen2.5-7B-Instruct, and Llama-3.1-8B-Instruct, harmful injections often succeed even when the pre-injection hidden state has high cosine similarity with the refusal prototype. For Mistral, the highest ASR coincides with the highest refusal-prototype similarity bin, summarized as “refused but didn’t resist.” PCA visualizations show that injecting harmful tokens at step 6 causes a sharp divergence from the refusal region to the accept region exactly at the injection step (Park et al., 3 Jun 2026).

To address this, the paper aligns generation trajectories rather than only states. It detects entry into refusal or compliance mode using cosine thresholds $y=(y_1,\dots,y_T)$5, injects the opposite-direction short phrase at the first triggering step, filters augmented continuations with Llama-Guard-3-8B, and trains on paired safe and unsafe trajectories using a reference-free SimPO objective. Iterative refinement recomputes the prototypes on the updated model, re-augments trajectories, and fine-tunes again (Park et al., 3 Jun 2026).

On AdvBench injection defense, the trajectory-aligned models substantially reduce ASR. For Llama-3.1-8B, “Injection, no defense” gives OM 92.12 and LG 89.70, while the aligned model gives OM 4.42 and LG 0.19. For Mistral-7B, the corresponding values are 85.77/82.49 and 0.00/0.96; for Qwen2.5-7B, 88.08/87.73 and 11.73/5.19. The method also generalizes to early-token exploits such as Prefilling, I-GCG, and PAIR, while MMLU and PROST remain within 1.1 points of baseline and XSTest over-refusal stays at or below 11.6% for Llama and Qwen, with Mistral dropping to 1.6% (Park et al., 3 Jun 2026).

This result sharply constrains a common interpretation of inference-time subspace alignment. A refusal-like hidden state is an informative diagnostic, but it is not a sufficient condition for robust refusal behavior. The paper’s formulation therefore shifts emphasis from occupying a refusal subspace to remaining aligned along the path of generation.

3. Sparse steering, stylistic subspaces, and preference subspaces

Recent LLM work also uses inference-time alignment to denote sparse or gated intervention at selected decoding steps. “Inference-time Alignment via Sparse Junction Steering” formulates a KL-regularized decoding problem and introduces a sparse intervention set $y=(y_1,\dots,y_T)$6, with steering applied only at high-entropy “junctions.” In its default implementation, SIA is a logit-space method rather than a hidden-state editing method:

$y=(y_1,\dots,y_T)$7

Absolute entropy gating dominates random, position, and attention gates; $y=(y_1,\dots,y_T)$8 works robustly across harmlessness, helpfulness, and honesty; and steering only 20% to 80% of tokens achieves superior alignment-efficiency trade-offs. For strong base models such as Qwen3, intervening on as few as 20% of tokens matches or even surpasses heavily post-trained instruct models, while computational cost is reduced by up to $y=(y_1,\dots,y_T)$9. The paper does describe hidden-state steering with a steering subspace $x$0 and update $x$1, but explicitly identifies this as an extension rather than the default implementation (Hu et al., 30 Jan 2026).

“Nudging: Inference-time Alignment of LLMs via Guided Decoding” uses a different subspace notion. Its claim is that alignment primarily alters behavior on a small subset of stylistic/discourse tokens, such as “Sure,” “Great,” “So,” “Let’s,” list markers, and transition words. The method gates intervention using the base model’s top-1 probability, with

$x$2

and inserts a short completion from a smaller aligned model. Empirically, when $x$3, the base and aligned top-1 tokens differ more than 90% of the time; setting the threshold to 0.5 captures more than 90% of alignment-related positions while fewer than 15% of positions fall below the threshold. Nudging typically modifies around 10% of tokens, and as low as about 5% in Gemma-2 can recover about 93% of large-chat average performance. Safety on just-eval safe reaches 5.00 for Llama-2 and 4.95 for Gemma-2 under GPT-4o judging (Fei et al., 2024).

“MAVIS: Multi-Objective Alignment via Value-Guided Inference-Time Search” introduces a further reinterpretation: the aligned object is a preference subspace in token-scoring space. For a fixed state $x$4, each objective defines a token-level value vector $x$5, and the guided logits become

$x$6

with user-specified weights $x$7. This yields an affine subspace of reachable guided logits. The paper reports that MAVIS generally matches or exceeds PPO rewards at similar KL in single-objective settings and matches or exceeds MORLHF’s Pareto front on HH-RLHF pairs while substantially outperforming RSoup and MOD across $x$8 sweeps (Carleton et al., 19 Aug 2025).

Across these papers, the phrase “subspace alignment” therefore covers several non-equivalent operations: sparse value-guided logit tilting, stylistic token anchoring, and user-weighted movement in a value span. This suggests that, in LLM inference, the word subspace often refers to a restricted control manifold in token or score space rather than only to a geometric basis in hidden-state space.

4. Classical geometric alignment and its test-time descendants

The most literal use of subspace alignment comes from domain adaptation. “Subspace Alignment For Domain Adaptation” represents source and target domains by PCA bases $x$9 and solves

$y_t \sim \pi(\cdot \mid x, y_{<t})$0

whose closed-form solution is

$y_t \sim \pi(\cdot \mid x, y_{<t})$1

Aligned source features are $y_t \sim \pi(\cdot \mid x, y_{<t})$2 and target features are $y_t \sim \pi(\cdot \mid x, y_{<t})$3. This formulation is unsupervised on the target side, admits principled subspace-dimension selection via an eigengap stability bound or MLE intrinsic dimension estimation, and was reported to outperform several earlier DA baselines on Office+Caltech-10, ILC-5, and ImageNet$y_t \sim \pi(\cdot \mid x, y_{<t})$4PASCAL-VOC-2007 (Fernando et al., 2014).

Later work adapts this idea to deep feature spaces and to test-time settings. “Revisiting Deep Subspace Alignment for Unsupervised Domain Adaptation” computes deep feature subspaces $y_t \sim \pi(\cdot \mid x, y_{<t})$5, aligns target to source with $y_t \sim \pi(\cdot \mid x, y_{<t})$6, and uses the ambient map $y_t \sim \pi(\cdot \mid x, y_{<t})$7 so that aligned target features are $y_t \sim \pi(\cdot \mid x, y_{<t})$8. The framework separates feature learning and distribution alignment, uses a primary–auxiliary optimization strategy, and explicitly supports progressive adaptation to new target domains at test-time, without requiring retraining of the model from scratch. The paper reports 90.02% average on ImageCLEF, 96.4% average on Digits, 73.02% on VisDA-2017, 65.1% average on Office-Home, 65.64% average on Office-Home PDA, and about a 6% average gain for progressive test-time adaptation relative to the originally deployed model without test-time adaptation (Thopalli et al., 2022).

“Test-time Adaptation for Regression by Subspace Alignment” argues that naive full-space feature alignment is ineffective in regression because features are distributed in a small output-relevant subspace. Its Significant-subspace Alignment (SSA) first performs PCA on source features, then aligns source and target statistics only in the top-$y_t \sim \pi(\cdot \mid x, y_{<t})$9 subspace, with dimension weights

$\mathcal{T}=\{(h_t,y_t)\}_{t=1}^T$0

where $\mathcal{T}=\{(h_t,y_t)\}_{t=1}^T$1 is the linear regressor and $\mathcal{T}=\{(h_t,y_t)\}_{t=1}^T$2 is the $\mathcal{T}=\{(h_t,y_t)\}_{t=1}^T$3-th source principal direction. The symmetric KL objective is computed only in this subspace, and only normalization affine parameters are adapted online. The paper reports, for example, SVHN$\mathcal{T}=\{(h_t,y_t)\}_{t=1}^T$4MNIST improvement from source $\mathcal{T}=\{(h_t,y_t)\}_{t=1}^T$5 to SSA $\mathcal{T}=\{(h_t,y_t)\}_{t=1}^T$6, California Housing improvement from $\mathcal{T}=\{(h_t,y_t)\}_{t=1}^T$7 to $\mathcal{T}=\{(h_t,y_t)\}_{t=1}^T$8, and highest mean $\mathcal{T}=\{(h_t,y_t)\}_{t=1}^T$9 on UTKFace corruptions (Adachi et al., 2024).

“An Unsupervised Tensor-Based Domain Alignment” extends the same family to multiway data. It aligns source and target tensors in an invariant subspace using mode-wise alignment matrices on the oblique manifold, regularized by variance-preservation terms. This tensor formulation generalizes existing tensor-based DA methods as special cases, and the paper reports that the oblique variant converges faster and achieves higher accuracy than Stiefel-constrained variants on both image and audio benchmarks (Lee et al., 26 Jan 2026).

These works preserve the original geometric meaning of the term: source and target data are mapped into a common or aligned subspace, often with explicit linear algebra and often with unlabeled target data available only at adaptation time.

5. Multimodal, audio-language, and federated formulations

In multimodal few-shot adaptation, “Constraining to Generalize: Subspace Tuning for Few-shot Generalization of Audio-LLMs” attributes the base-to-new trade-off to zero-shot drift in the text embedding space and proposes Subspace Tuning (SubT). Base prototypes are factorized as

$h_t$0

and adaptation learns only a shared basis $h_t$1, yielding

$h_t$2

At inference on unseen classes, the learned basis shift is transferred by

$h_t$3

and then modulated by subspace-aware gating

$h_t$4

Across 11 audio benchmarks, the reported averages are Base 87.89, New 62.49, H 71.79 for SubT, and Base 87.89, New 63.79, H 72.52 for SubT† with gating, compared with Base 62.66, New 61.17, H 60.02 for zero-shot and Base 77.08, New 59.98, H 65.27 for CLIP-Adapter (Jang et al., 17 Jun 2026).

A different modality of inference-time subspace alignment appears in vertical federated learning. “X-VFL: A New Vertical Federated Learning Framework with Cross Completion and Decision Subspace Alignment” defines the decision subspace as the representation space induced by the top model $h_t$5. For two clients, DS-Align imposes both reconstruction alignment,

$h_t$6

and local-to-global alignment,

$h_t$7

Because the global embedding is an average rather than a concatenation, the top model can be reused for single-client inference, and $h_t$8 forces $h_t$9. The paper reports about a 15% improvement in accuracy on CIFAR-10 and about a 43% improvement on MIMIC-III, together with convergence rates of $(x,y_{<t})$0 for SGD-type algorithms and $(x,y_{<t})$1 for PAGE-type algorithms (Yao et al., 7 Aug 2025).

These multimodal and federated variants preserve the core idea of constrained inference-time adaptation, but the aligned object is no longer necessarily a classical PCA subspace. It may instead be a learned basis in text-prototype space or the decision space induced by a task head.

6. Diffusion guidance, limitations, and open directions

“Inference-Time Alignment for Diffusion Models via Doob’s Matching” introduces a principled diffusion analogue. A pretrained base score model is tilted at sampling time toward a target distribution $(x,y_{<t})$2 through a Doob $(x,y_{<t})$3-transform. The guided reverse-time SDE is

$(x,y_{<t})$4

For subspace alignment, the target tilt is restricted to a feature map $(x,y_{<t})$5 by setting $(x,y_{<t})$6 and $(x,y_{<t})$7, which yields

$(x,y_{<t})$8

The paper estimates the guidance by gradient-penalized regression and proves non-asymptotic convergence of the learned guidance together with $(x,y_{<t})$9-Wasserstein guarantees for the resulting generated distribution under bounded-support and bounded-weight assumptions (Chang et al., 10 Jan 2026).

A persistent theme across the literature is that alignment in a subspace is rarely sufficient by itself. In LLM safety, high cosine similarity to a refusal direction does not imply robustness under perturbation; path dependence dominates (Park et al., 3 Jun 2026). In sparse steering, excessive intervention or misidentified junctions can induce drift or “entropy rebound,” and value-model quality becomes the bottleneck (Hu et al., 30 Jan 2026). In guided decoding, safety inherits from the smaller nudging model, and oversteering can cause repetition or unnatural phrasing (Fei et al., 2024). In regression TTA, choosing $h_{t+1}=f_\theta(x,y_{\le t})$0 above the effective subspace rank can destabilize the loss (Adachi et al., 2024). In ALM few-shot transfer, weakly aligned unseen classes motivate gating because ungated transfer can cause negative transfer (Jang et al., 17 Jun 2026).

Several open problems recur. One is richer geometry: refusal cones or multi-dimensional refusal subspaces beyond a single direction in LLM safety, and alternative alignment metrics such as principal angles or CCA-style projections in multimodal transfer (Park et al., 3 Jun 2026, Jang et al., 17 Jun 2026). Another is stochastic or search-based decoding, since some safety evaluations use greedy decoding and robustness under sampling remains unresolved (Park et al., 3 Jun 2026). A third is formal robustness guarantees: diffusion guidance now has non-asymptotic theory, but corresponding guarantees for autoregressive generation under token injection remain largely open (Chang et al., 10 Jan 2026, Park et al., 3 Jun 2026).

The literature therefore supports a precise but non-monolithic understanding of inference-time subspace alignment. In its classical form, it is source–target geometric alignment in a low-dimensional basis. In contemporary generative modeling, it also names a broader set of inference-time controls that act through refusal directions, stylistic tokens, value spans, decision spaces, or feature-conditioned guidance. The unifying idea is not a single mathematical object, but a shared strategy: alignment is imposed through a restricted structure during inference, with the aim of preserving the base model’s broader capabilities while selectively correcting behavior.