RouteSAE: Sparse Routing for LLMs & Road Safety

• RouteSAE is a dual-framework approach that uses a shared TopK sparse autoencoder with dynamic routing to extract monosemantic features from transformer layers in large language models.
• It extends to roadway safety by applying task-specific attention to panoramic images, enabling multi-task prediction and efficient deep neural assessment of street-level safety.
• Empirical results show a 22% improvement in both feature interpretability and reconstruction error over conventional per-layer methods, underlining its cross-domain effectiveness.

RouteSAE refers to two distinct frameworks addressing different domains, both unified by the principle of route selection via sparse or attention-based mechanisms: (1) the Route Sparse Autoencoder (RouteSAE) for scalable mechanistic interpretability of LLMs (Shi et al., 11 Mar 2025), and (2) a route-aware architecture for automated roadway safety assessment based on deep neural networks, following the FARSA methodology (Song et al., 2019). Both share a focus on efficient, interpretable feature extraction across multiple layers or tasks.

1. Mechanistic Interpretability in LLMs via RouteSAE

Overview

RouteSAE in the LLM context is a framework designed to extract monosemantic and interpretable features from all depths of a transformer-based model. It combines a single shared TopK sparse autoencoder (SAE) with a routing network that dynamically identifies the most informative residual stream layer to encode, yielding considerable parameter savings and enhanced interpretability compared to single-layer or per-layer baselines (Shi et al., 11 Mar 2025).

2. RouteSAE Architecture

Routing Mechanism

Let residual stream activations at $L$ layers be $\mathbf{x}_i\in\mathbb{R}^d$ for $i=0,\dots,L-1$. RouteSAE first computes a pooled vector,

$\mathbf{v}=\sum_{i=0}^{L-1}\mathbf{x}_i,$

which is projected to logits via

$$\boldsymbol{\alpha}=W_{\mathrm{router}}\,\mathbf{v},$$

with $W_{\mathrm{router}}\in\mathbb{R}^{L\times d}$. The softmax normalization produces layer selection probabilities $p_i$, and in hard-routing (main variant), the routed input is

$$\mathbf{x}_{\mathrm{route}} = p_{i^*}\mathbf{x}_{i^*},\quad i^* = \arg\max_i p_i.$$

This vector is then processed through a single shared TopK SAE across all layers.

Sparse Autoencoder

The SAE encoder and decoder are defined by weights $W_{\mathrm{enc}}\in\mathbb{R}^{M\times d}$ and $W_{\mathrm{dec}}\in\mathbb{R}^{d\times M}$, with pre-activation bias $\mathbf{b}_{\mathrm{pre}}$. Latent codes are computed by

$$\mathbf{z}=\mathrm{TopK}\big(W_{\mathrm{enc}}(\mathbf{x}_{\mathrm{route}}-\mathbf{b}_{\mathrm{pre}})\big)\in\mathbb{R}^M,$$

where only the largest $K$ pre-activations are retained, enforcing strict sparsity. Reconstruction is

$$\hat{\mathbf{x}}_{\mathrm{route}}=W_{\mathrm{dec}}\,\mathbf{z}+\mathbf{b}_{\mathrm{pre}}.$$

Training jointly minimizes the reconstruction loss

$$\mathcal{L}_{\mathrm{RouteSAE}} = \lVert \mathbf{x}_{\mathrm{route}} - \hat{\mathbf{x}}_{\mathrm{route}}\rVert_2^2,$$

with unit-norm regularization on decoder columns.

Parameter Efficiency

By sharing the SAE across all routed layers, RouteSAE adds only $\mathcal{O}(Ld)$ parameters (router matrix and bias) rather than $\mathcal{O}(LMd)$ as in per-layer approaches like Crosscoder. Empirically, this yields an approximately $8\times$ reduction in parameter overhead for comparable or superior interpretability and reconstruction fidelity (Shi et al., 11 Mar 2025).

3. Interpretability and Performance Metrics

Interpretability is quantified along two axes:

• Interpretable Feature Count: Features are retained if they activate on at least four distinct high-activation contexts across a validation corpus. Features passing this threshold are deemed interpretable at the given threshold.
• Interpretability Score: For $N=100$ sampled features, GPT-4o assigns an “interpretation category” (low-level, high-level, or undiscernible) and a monosemanticity score $s_j\in[1,5]$. The mean interpretability score is $\frac{1}{N}\sum_{j=1}^{N} s_j$.

RouteSAE achieves 22.5% more interpretable features and a 22.3% higher interpretability score relative to a single-layer TopK SAE at matched sparsity ($K=64$), demonstrating its advantage in both feature richness and semantic purity. In downstream KL evaluation (residual replacement task), RouteSAE establishes a superior sparsity–KL Pareto frontier and outperforms Crosscoder in normalized reconstruction error (0.18 vs. 0.35). Ablation removing routing collapses gains to single-layer SAE baseline levels (Shi et al., 11 Mar 2025).

4. Use Cases, Applications, and Extension Strategies

Interpretability and Feature Manipulation

Because RouteSAE produces a single, aligned feature space spanning all routed layers, it enables:

• Cross-layer feature discovery—encompassing both early polysemy disambiguation and deep pattern integration in one model.
• Targeted interventions—mechanistically altering generation by clamping or rescaling individual features at inference.

Proposed Extensions

• Router augmentation with multi-head or attention-based aggregation.
• Routing regularizers for richer soft-routing mixtures.
• Integration into larger LLMs, Mixture-of-Experts, or encoder–decoder architectures.
• Use of proximal/alternating-minimization schemes to further minimize reconstruction error while preserving sparsity.

5. RouteSAE for Automated Roadway Safety Assessment

Following the FARSA methodology, RouteSAE also denotes a deep neural pipeline for visual safety rating of street-level panoramas (Song et al., 2019):

• Input: $x\in\mathbb{R}^{224\times960\times3}$ panoramas (Google Street View, GSV).
• Outputs: Primary star rating $y^s\in\{1,\dots,5\}$ (usRAP standard, one-hot or simplex $\hat{y}^s\in\Delta^4$) and $M=17$ auxiliary discrete roadway attributes (e.g., median type, intersection channelization).
• Backbone Architecture: Truncated VGG-16 (conv1–conv5), followed by $1\times1$ ReLU conv, yielding $S_2\in\mathbb{R}^{7\times30\times512}$, then flattened to $\widetilde{S}_2\in\mathbb{R}^{210\times512}$.
• Task-specific Attention: For each task $t$, attention weights $a_t\in\mathbb{R}^{210}$, softmax-normalized, yield a fused feature $$f_t=\sum_{i=1}^{210}\alpha_{t,i}\widetilde{S}_{2,i}\in\mathbb{R}^{512}$$. These feed task-specific classification heads.

Multi-task and Semi-supervised Training

The total loss is a weighted sum: $$\mathcal{L}_{\rm total} = \lambda_s \mathcal{L}_s + \lambda_m \mathcal{L}_m + \lambda_u \mathcal{L}_u,$$ combining a star-rating loss (classification + ordinal regression), multi-task cross-entropy for auxiliary attributes, and an unsupervised geographic consistency loss for adjacent panorama pairs. Semi-supervised batches (16 labeled, 16 unlabeled) encourage consistency and regularization.

Preprocessing, Training, and Evaluation

Panoramas are preprocessed by orientation, cropping, resizing, and stratified test split spatially separated by at least 300 m from training. Training uses Adam optimizer and $L_2$ weight decay. The best model achieves a macro average top-1 test accuracy of 46.91% on star rating, with attention, multi-task, and unsupervised learning yielding cumulative gains over the backbone alone.

Star-rating confusion concentrates on rare 1-star roads due to data imbalance; auxiliary attribute accuracy exceeds random priors across tasks.

6. Limitations, Contingencies, and Directions for Further Research

LLM Interpretability RouteSAE

• Absence of routing regularizer or multi-head attention in the standard variant may limit the diversity of routed mixtures. A plausible implication is that more sophisticated router designs could further enhance cross-layer feature compositionality.
• Performance depends on the choice of routed layers; inappropriate selection collapses to single-layer behavior.
• The small parameter overhead is contingent on the number of routed layers $L$ and hidden size $d$; extremely large $L$ may impose non-negligible costs.

Roadway Safety RouteSAE

• The main performance bottleneck is label scarcity for rare, high-risk (1-star) roads.
• Geographic and task-specific pooling or satellite imagery integration are proposed to improve robustness where street-level features are occluded or ambiguous.
• Potential operational deployment involves incorporating model predictions into routing engines by adjusting edge costs as $c_{e}'=c_e\,\phi(\hat{y}_e^s)$, highly penalizing risky segments.

7. Summary Table of Key RouteSAE Mechanisms

Context	Routing Source	Feature Aggregation	Primary Gains
LLM interpretability (Shi et al., 11 Mar 2025)	Residual stream layers	Shared TopK SAE	Multi-layer feature extraction, parameter efficiency, high interpretability
Road safety (Song et al., 2019)	Panoramic image regions	Task-specific attention	Multi-task/attribute prediction, label-efficient regularization

RouteSAE thereby denotes a set of architectures unifying sparse- or attention-based routing with shared feature extraction, offering scalable interpretability and multi-layer/multi-task efficiency in both LLMs and computer vision safety assessment contexts (Shi et al., 11 Mar 2025, Song et al., 2019).