Domain Restriction via Multi SAE Layer Transitions

Abstract: The general-purpose nature of LLMs presents a significant challenge for domain-specific applications, often leading to out-of-domain (OOD) interactions that undermine the provider's intent. Existing methods for detecting such scenarios treat the LLM as an uninterpretable black box and overlook the internal processing of inputs. In this work we show that layer transitions provide a promising avenue for extracting domain-specific signature. Specifically, we present several lightweight ways of learning on internal dynamics encoded using a sparse autoencoder (SAE) that exhibit great capability in distinguishing OOD texts. Building on top of SAEs representation transitions enables us to better interpret the LLM internal evolution of input processing and shed light on its decisions. We provide a comprehensive analysis of the method and benchmark it with the gemma-2 2B and 9B models. Our results emphasize the efficacy of the internal process in capturing fine-grained input-related details.

• The paper introduces a novel sequential modeling approach using multi-layer sparse autoencoder transitions for precise domain restriction in LLMs.
• It leverages global density filtering and Top-K binarization to generate sparse, interpretable feature codes across transformer layers.
• Experimental results show robust performance with a Markov model, validating its efficiency in detecting both far- and near-OOD cases.

Domain Restriction via Multi SAE Layer Transitions: Technical Analysis

Introduction and Motivation

This paper introduces a methodology for domain restriction in LLMs via sequential modeling of internal representation transitions, specifically leveraging sparse autoencoders (SAEs) to decompose residual stream activations at multiple layers. The method is motivated by the need for high-precision scope gating in practical LLM deployments, particularly when routing requests to domain-specific agentic workflows. Existing OOD detectors for LLMs operate mainly as output-level black boxes, lacking both transparency and alignment with internal representations. By contrast, the approach in this work exploits depthwise transitions in sparse, interpretable SAE feature space, learned on in-domain data only, to provide fine-grained, light-weight, and interpretable OOD discrimination without requiring labeled OOD data or fine-tuning of the host model.

Methodology

Sparse Autoencoder-Based Representation

The core technique is to utilize pre-trained, layer-wise SAEs to decompose the high-dimensional, superposed residual-stream activations into a sparse set of monosemantic features. Each layer thus yields a sparse feature code per input token, which is then pooled across tokens to form a trajectory of sparse distributed representations (SDRs) across network depth.

Key properties:

• Global density filtering: Common (high-frequency) features are masked to improve domain specificity.
• Top-K binarization: Only the most salient $k$ features per layer are retained, supporting transition modeling.
• SDR trajectories: Inputs are represented as sequences of binarized, sparse feature sets across transformer layers.

Sequential Anomaly Scoring

A lightweight sequential model is trained, using only in-domain data, to assign an anomaly score to the depthwise SDR trajectory of new inputs. Several scoring architectures are evaluated:

• First-order Markov transition model: Estimates the probability of adjacent-layer feature transitions.
• HTM (Hierarchical Temporal Memory): Models higher-order sequence dependencies.
• RNN (LSTM/GRU): Predicts next-layer activation patterns.

The Markov backend—despite its simplicity—exhibits robust performance, reinforcing the hypothesis that domain specificity is primarily encoded in local layer-to-layer transitions rather than in higher-order dependencies.

Experimental Results

Datasets and Benchmarks

Experiments use Gemma2-2B and Gemma2-9B models. Evaluation covers far-OOD scenarios (e.g., train on 20 Newsgroups, test on SST-2, MNLI, RTE, IMDB, CLINC150) and near-OOD scenarios (semantics-separated splits of AGNews, ROSTD, SNIPS, CLINC150). The baseline comparator is the likelihood-ratio (LR) detector between base and in-domain-finetuned models [see also Zhang et al., 2025].

Key Numerical Results

Far-OOD

• Gemma2-9B + 16k_SAE + Markov: AUROC up to 0.99, AUPR up to 0.98, FPR@95% sometimes below 0.07.
• Baseline LR: Remains superior but requires two models; this method is more lightweight.

Near-OOD

• On AGNews and SNIPS, the proposed approach matches or approaches the LR baseline in AUROC and FPR@95.
• On CLINC150 (fine-grained multi-intent), performance degrades (e.g., AUROC ≈ 0.74 to 0.83 for 2B model), attributed to representation resolution mismatch.

Ablation

• Markov transition modeling matches or outperforms more complex sequence models (HTM, RNN).
• Performance collapses when operating directly on raw activations instead of SAE-derived features (AUROC ≈ 0.75), validating the need for disentangled feature representations.

Representation Dynamics

• In-domain consistency peaks in mid-depth layers ("semantic trunk"), with stable Top-K feature overlap.
• Trajectory-based (i.e., sequential) signatures of SDR transitions provide stronger and more robust domain cues than static overlap.
• Registry-based explicit modeling of feature-tuple trajectories is less effective than learned scoring functions due to noise accumulation.

Interpretability and Case Studies

The approach supports mechanistic interpretability via feature-level transitions. Examination of errors (e.g., "hard-OOD" false negatives) reveals that misclassified OOD samples have content that semantically overlaps with the in-domain class, and their internal transition dynamics closely match those of in-domain samples. Inspection of feature transitions via Neuropedia decoding further substantiates these findings.

Discussion and Implications

Theoretical Implications

• This work operationalizes a "features not neurons" paradigm in mechanistic interpretability, illustrating that sparse, compositional feature transitions capture domain-specific processing in LLMs more robustly than raw activations.
• It demonstrates that domain boundaries in internal representation space can be modeled as sequential patterns rather than static clusters.
• The Markovian nature of effective discriminative cues across transformer depth is empirically validated.

Practical Implications

• The proposed gating mechanism is efficient, resource-light, and does not require OOD samples or multi-model inference.
• The method is well suited for real-world agent-based pipelines, offering more interpretable "why" outputs for scope gating decisions.
• Fine-grained task differentiation (as in CLINC150) challenges the approach, exposing the limit of current SAE dictionary expressivity.

Future Directions

Several promising avenues are identified:

• Expanding SAE expressivity (e.g., more features, better dictionaries) to match the granularity required by complex intent taxonomies.
• Incorporating graded feature strength, not just binarized activation, for improved discrimination.
• Integrating dynamic, context-adaptive sparsity for improved robustness.
• Application to tool-based and retrieval-augmented LLM agent routing.

Conclusion

This work establishes the viability of domain restriction for LLMs via sequential modeling of SAE-derived layer transition patterns. The approach offers a powerful scope gating signal in data regimes where OOD data is scarce or unavailable, and supports interpretation via feature-level analysis. While effective in coarse- and moderate-grained domain splits, its resolution is currently limited by the underlying SAE feature dictionary, indicating fruitful directions for further research in mechanistic interpretability, dictionary learning, and the tailoring of LLM internals for agentic modularity.