---
title: Inference-Time Subspace Alignment
url: https://www.emergentmind.com/topics/inference-time-subspace-alignment
type: topic
---

# Inference-Time Subspace Alignment

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 [2606.04778] [2602.21215] [2410.09300] [2601.06514] [1409.5241].

## 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 [2606.04778].

## 2. Trajectory-level alignment in autoregressive language models

A safety-aligned autoregressive language model $\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 [2606.04778].

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” [2606.04778].

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 $\bar{h}_r$ and $\bar{h}_a$ are the mean hidden states over refusal and compliance samples, the refusal prototype is
$$
p_r=\bar{h}_r-\bar{h}_a,\qquad p_a=-p_r,
$$
and the alignment measure at step $t$ is
$$
H_{i,t}=\cos(h_{i,t},p_r).
$$
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 [2606.04778].

To address this, the paper aligns **generation trajectories** rather than only states. It detects entry into refusal or compliance mode using cosine thresholds $\tau_r,\tau_a$, 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 [2606.04778].

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% [2606.04778].

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 $S \subseteq \{1,\dots,T\}$, 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:
$$
\pi_{\text{sparse}}(y_t\mid s_t)=
\begin{cases}
\frac{1}{Z_t}\pi_{\text{base}}(y_t\mid s_t)\exp\!\big(\beta\,\widehat{V}(s_{t+1})\big), & t\in S,\\
\pi_{\text{base}}(y_t\mid s_t), & t\notin S.
\end{cases}
$$
Absolute entropy gating dominates random, position, and attention gates; $\tau \approx 1.0$ 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 $6\times$. The paper does describe hidden-state steering with a steering subspace $U$ and update $h'_t=h_t+U\gamma_t$, but explicitly identifies this as an extension rather than the default implementation [2602.21215].

“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
$$
g_t=\mathbf{1}\!\big[\text{top1}(P_{\text{base}}(\cdot\mid q,a_{<t}))<\gamma\big],
$$
and inserts a short completion from a smaller aligned model. Empirically, when $\text{top1}(p_{\text{base}})<0.1$, 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 [2410.09300].

“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 $s$, each objective defines a token-level value vector $G_m(s)=Q_m^*(s,\cdot)$, and the guided logits become
$$
L_{\boldsymbol{\lambda},\beta}(s)=L_0(s)+\beta\sum_{m=1}^M \lambda_m G_m(s),
$$
with user-specified weights $\lambda_m$. 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 $\lambda$ sweeps [2508.13415].

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 $A,B\in\mathbb{R}^{D\times k}$ and solves
$$
\min_M \|AM-B\|_F^2,
$$
whose closed-form solution is
$$
M^*=A^\top B.
$$
Aligned source features are $Z_s=X_sAM^*=X_s(AA^\top B)$ and target features are $Z_t=X_tB$. 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$\rightarrow$PASCAL-VOC-2007 [1409.5241].

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 $U_s,U_t$, aligns target to source with $\Phi^*=U_t^\top U_s$, and uses the ambient map $M^*=U_t\Phi^*U_s^\top=(U_tU_t^\top)(U_sU_s^\top)$ so that aligned target features are $z_t'=z_tU_t\Phi^*U_s^\top$. 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 [2201.01806].

“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-$K$ subspace, with dimension weights
$$
\alpha_d=1+\left|w^\top v_d^s\right|,
$$
where $w$ is the linear regressor and $v_d^s$ is the $d$-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$\rightarrow$MNIST improvement from source $R^2=0.406$ to SSA $R^2=0.511\pm0.03$, California Housing improvement from $R^2=0.605$ to $0.639\pm0.00$, and highest mean $R^2=0.731$ on UTKFace corruptions [2410.03263].

“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 [2601.18564].

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-Language Models” 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
$$
F_{\text{base}}=U\Sigma V_0^\top,\qquad C:=U\Sigma,
$$
and adaptation learns only a shared basis $V_{\text{ft}}^\top$, yielding
$$
\hat{F}_{\text{base}}=\text{Norm}(CV_{\text{ft}}^\top),
\qquad
F_{\text{base}}^{\text{tuned}}=\text{Norm}(F_{\text{base}}+\hat{F}_{\text{base}}).
$$
At inference on unseen classes, the learned basis shift is transferred by
$$
T:=V_0V_{\text{ft}}^\top,\qquad \hat{F}_{\text{new}}=\text{Norm}(F_{\text{new}}T),
$$
and then modulated by **subspace-aware gating**
$$
\beta_i:=\|f_i^{\text{new}}V_0\|_2,\qquad
f_i^{\text{new,final}}=\text{Norm}(f_i^{\text{new}}+\beta_i\hat{f}_i^{\text{new}}).
$$
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 [2606.18560].

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$. For two clients, DS-Align imposes both reconstruction alignment,
$$
L_{\text{DSAlign}_1}
=
\ell(h(\tilde{E}_a),h(E_a))+\ell(h(\tilde{E}_b),h(E_b)),
$$
and local-to-global alignment,
$$
L_{\text{DSAlign}_2}
=
\ell\!\left(h(E_a),h\!\left(\frac{E_a+E_b}{2}\right)\right)
+
\ell\!\left(h(E_b),h\!\left(\frac{E_a+E_b}{2}\right)\right).
$$
Because the global embedding is an average rather than a concatenation, the top model can be reused for single-client inference, and $L_{\text{DSAlign}_2}$ forces $h(E_i)\approx h(\text{Avg}(E_1,\dots,E_k))$. 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 $O(1/\sqrt{T})$ for SGD-type algorithms and $O(1/T)$ for PAGE-type algorithms [2508.05568].

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 $q_0(x)\propto w(x)p_0(x)$ through a Doob $h$-transform. The guided reverse-time SDE is
$$
dZ_t^\leftarrow=
\big[Z_t^\leftarrow+2\nabla\log p_{T-t}(Z_t^\leftarrow)+2\nabla\log h^*(t,Z_t^\leftarrow)\big]dt+\sqrt{2}\,dB_t.
$$
For **subspace alignment**, the target tilt is restricted to a feature map $\Phi:\mathbb{R}^d\to\mathbb{R}^k$ by setting $w(x)=w_{\text{sub}}(\Phi(x))$ and $h(t,x)=h_{\text{sub}}(t,\Phi(x))$, which yields
$$
\nabla_x\log h(t,x)=J_\Phi(x)^\top \nabla_z \log h_{\text{sub}}(t,z).
$$
The paper estimates the guidance by gradient-penalized regression and proves non-asymptotic convergence of the learned guidance together with $2$-Wasserstein guarantees for the resulting generated distribution under bounded-support and bounded-weight assumptions [2601.06514].

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 [2606.04778]. In sparse steering, excessive intervention or misidentified junctions can induce drift or “entropy rebound,” and value-model quality becomes the bottleneck [2602.21215]. In guided decoding, safety inherits from the smaller nudging model, and oversteering can cause repetition or unnatural phrasing [2410.09300]. In regression TTA, choosing $K$ above the effective subspace rank can destabilize the loss [2410.03263]. In ALM few-shot transfer, weakly aligned unseen classes motivate gating because ungated transfer can cause negative transfer [2606.18560].

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 [2606.04778] [2606.18560]. Another is **stochastic or search-based decoding**, since some safety evaluations use greedy decoding and robustness under sampling remains unresolved [2606.04778]. 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 [2601.06514] [2606.04778].

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.

Source: https://www.emergentmind.com/topics/inference-time-subspace-alignment