Learning Cross-Model Activation Alignments with Explicit Many-to-Many Layer Maps

Motivation

• As open-weight model releases accelerate, rebuilding probes, steering vectors, and other activation-based tools per model increases costs and complicates comparison across models.The tools operate on a target model’s hidden activations and are not directly applicable to successor or competitor models.Source: §1 · p. 1
• Activation alignment maps one model’s layerwise activations to another’s, enabling applications developed for one model to be reused without retraining or modifying either model.It can also reveal structural relationships, including layer correspondences.Source: §1 · p. 2

Prior work and gap

• Prior work commonly fixes a one-to-one layer mapping independent of the input and model pair, while extending flexible mappings to input-dependent settings remains unclear.Solving a separate optimization problem for each input would also be required by such an extension.Source: §1 · p. 2

Our approach

• MATCHA trains a lightweight cross-attention network to predict an interpretable layer map from the input, avoiding separate per-input optimization.Target-layer queries attend to input-dependent source-layer keys; a shared projection and feed-forward network transform features, and both mappings are trained jointly.Source: §1 · pp. 2–3

Main findings and contributions

• Across 42 pairs of seven models, MATCHA reconstructs target activations more accurately than all baselines and achieves stronger retrieval metrics.Its layer maps are broadly monotone but many-to-many, and it improves transfer of probes and steering interventions across four datasets.Source: §1 · p. 3

Cross-model activation maps

• Cross-model activation maps commonly impose fixed one-to-one layer pairings, while newer methods estimate softer many-to-many correspondences across layers or features.Examples include relative-depth or terminal pairing, cross-coders, and Hierarchical Optimal Transport.Source: §2 · p. 3

Applications of activation maps

• Activation maps have been used to transfer steering interventions, key–value caches, mechanistic circuits, and representations for classification and detection.Other applications include communication between language-model agents and guiding cross-architecture weight-space merging.Source: §2 · p. 3

Evaluating cross-model representations

• Representation-similarity metrics compare space geometry without constructing an explicit map, whereas model stitching tests functional compatibility by injecting transformed representations into another network.This paper evaluates mappings through retrieval, reconstruction, and downstream behavior.Source: §2 · pp. 3–4

Problem formulation

• For each prompt, the source and target activation stacks contain final-token hidden states from every transformer block.The models may have different layer counts and hidden dimensionalities.Source: §3 · p. 4
• The goal is to learn an activation alignment map that maps source activations to the target activations produced by the same prompt.The mapped source representation should approximate the target representation.Source: §3 · p. 4

4 MATCHA: Many-to-many Alignment Through Cross-attention over Hidden Activations

• MATCHA factorizes the alignment into layer and feature maps to make layer contributions and feature transformations explicit while retaining flexibility to subsume prior one-to-one mappings.The construction is guided by explicitness and flexibility as design principles.Source: §4 · pp. 4–5
• The alignment uses final-token states, which under causal masking are the only positions with access to the full prompt.The paper notes that this state is commonly used for representation readout.Source: §4 · p. 4
• The layer map returns an explicit target-by-source matrix whose entries give each source layer’s additive contribution to a target layer.Because the correspondence is materialized as a matrix, it can be read off directly.Source: §4 · p. 5
• Cross-attention uses learned, input-independent target-layer queries and source-activation-dependent keys; each attention weight assigns a source layer’s contribution to a target layer.The keys are formed by a learnable projection of the source activations.Source: §4 · p. 5
• Under mild conditions, the layer map can approximate prior constant one-to-one mappings and any constant soft many-to-many map.Source: §4 · p. 5
• The feature map carries layer-mixed source representations into the target hidden space using a layer-shared projection and residual feed-forward block.The feed-forward block makes the feature transformation fully expressive.Source: §4 · p. 5

5 Experiments

• The study evaluates alignments across seven instruction-tuned models, yielding 42 non-degenerate ordered pairs, and compares fixed one-to-one layer maps with soft many-to-many maps.For each pair, LIMA’s 1,030-conversation training split was divided into training and validation subsets; alignment quality was evaluated on the official test set of 300 held-out prompts.Source: §5 · p. 6

5.1 Evaluation of Cross-Model Alignment

• Around 0.12 R@1 gain in the smaller-to-larger setup; MATCHA performs best across metrics and model-size configurations.Relative Depth is the second-best overall method.Source: Table 1 · p. 7, §5.1 · p. 7
• Evaluation combines nearest-neighbor retrieval, token-neighbor overlap, and reconstruction MSE; higher retrieval scores and lower MSE indicate better performance.Source: §5.1 · p. 7

5.2 Learned Layer Correspondences

• Learned correspondences are generally many-to-many, with target layers combining multiple source layers; maps involving Gemma show less clear correspondence.Maps are broadly concentrated near relative depth, but can curve into a middle-layer “belly” when model depths differ.Source: §5.2 · p. 8, Figure 6 · p. 8

5.3 Downstream Applications

• Across three datasets, MATCHA outperforms other alignment techniques by a wide margin when transferring refusal-suppression steering.The setup maps target steering-induced changes back to source layers; the native adapter provides an upper bound and no steering a lower bound.Source: §5.3 · p. 9, Figure 8 · p. 9
• For TriviaQA hallucination detection, MATCHA outperforms all baselines, improving AUC by approximately three points over the second-best method in smaller-to-larger transfers.Relative Depth is the strongest baseline overall.Source: §5.3 · p. 9, Table 2 · p. 10

5.4 Ablation Studies

• Restricting each target layer to one source layer performs substantially worse than using multiple sources, and the full map beats top-k approximations.Source: §5.4 · pp. 9–10
• Removing input dependence or using an unstructured MLP degrades performance across metrics, supporting MATCHA’s input-dependent map and separate layer- and feature-map design.Source: §5.4 · p. 10, Table 3 · p. 10

• MATCHA applies the learned alignment in approximately 0.06 milliseconds and decomposes it into an interpretable cross-layer map and a feature transformation.Source: §6 · p. 10
• Across 42 ordered model pairs, MATCHA consistently outperforms prior baselines, with particularly strong gains for cross-family and smaller-to-larger transfers.Source: §6 · p. 10
• The learned layer maps are broadly depth-monotone yet inherently multi-layer, rather than one-to-one.Source: §6 · p. 10
• The alignments enable probes and steering interventions to transfer across models, supporting reuse of activation-based tools with new LLMs without retraining each tool.Source: §6 · p. 10

• Transferring steering interventions is dual-use and could be misused to weaken safeguards, so the authors recommend safety evaluation and appropriate access controls before deployment.Source: §7 · p. 11

Limitations

• The method requires hidden-state access, limiting use to open-source models or direct implementation by closed-source model providers; it also trains a separate map per pair.A joint map trained across multiple pairs is left for future work.Source: Limitations · p. 11

A.1 Activation Normalization

• Before training, layer-specific means and standard deviations are estimated on alignment-training data and used to normalize activations, preventing larger scales from dominating the objective.Outputs are denormalized with target-layer statistics for activation editing and evaluation, and the calibration statistics remain fixed across training, validation, and testing.Source: §A.1 · p. 16

A.2 Extended Metric Definitions

• Recall@1 checks whether the true target is among the target-space representations nearest to the mapped prediction, while Precision@1 queries from the true target toward the mapped space.Source: §A.2 · pp. 16–17
• Jaccard@5 compares source- and mapped-space neighbor sets by intersection over union, penalizing both lost neighbors and spurious additions.For K>1, the union penalty distinguishes Jaccard from Overlap.Source: §A.2 · p. 17, Figure 12 · p. 17
• Token-Jaccard applies the Jaccard measure to top-next-token sets, and is the only listed metric that depends on the target model's unembedding readout.Source: §A.2 · p. 17

B Extended Experimental Section

• The extended experimental section reports AdamW with zero weight decay, cosine scheduling, 10% linear warm-up, and early stopping with patience 5; maps use last-token residual-stream states from frozen models.The states are extracted at every layer using chat-templated prompts and one forward pass.Source: B Extended Experimental Section · p. 19

B.1 Models Description

• The model-pair groups contain 18 bigger-to-smaller pairs, 18 smaller-to-bigger pairs, and 6 same-size directional pairs.Source: §B.1 · p. 19

B.2 Dataset Description; B.3 Additional Results

• LIMA prompts supply disjoint alignment training, validation, and test activations, without using task labels or responses to train the maps.The additional-results variants remove the feed-forward component or use a fixed average map computed on held-out validation data.Source: §B.2 · p. 19, §B.3 · p. 20

B.2 Dataset Description

• TriviaQA questions and reference answers provide binary correctness labels for training and evaluating layerwise correctness probes.Source: §B.2 · p. 19
• AdvBench, HarmBench, and JailbreakBench provide three held-out benchmarks for evaluating refusal-suppression steering.HarmBench covers harmful behaviors and misuse categories; JailbreakBench contains 100 annotated misuse behaviors.Source: §B.2 · pp. 19–20
• Refusal-direction prompts construct per-layer contrastive steering directions and select intervention layers and coefficients, with prompts kept separate from final evaluation.Source: §B.2 · p. 20

B.3.1 Retrival-Based Metrics per Model Pair

• MATCHA’s per-pair retrieval results are reported for Recall@1, Precision@1, Jaccard@5, and Token-Jaccard@5, grouped by source model.The supplied passages identify full per-pair tables with seed means and standard deviations but do not expose their cell values.Source: §B.3.1 · p. 20, Table 14 · p. 30, Table 15 · p. 31

B.3.2 MSE per Model Pair

• The MSE appendix reports full per-pair loss results, grouped by source model, with seed means and standard deviations.The supplied passages do not provide the table’s cell values.Source: §B.3.2 · p. 34, Table 18 · p. 35

B.3.3 Probing

• Probing results compare the receiver’s own probe with a receiver probe re-selected by validation on the same transferred representation, and report the validation-winning receiver layer for each method and pair.Receiver depth ranges from 16 to 36 layers, and the selected layer depends on the method and pair.Source: §B.3.3 · pp. 36, 39, Table 19 · p. 37, Table 20 · p. 38

B.3.4 Refusal-Suppression Steering

• Transferred refusal steering maps source activations into target space, edits the selected target layer, maps the edit back, and applies source-layer updates in a second forward pass.The alignment maps are frozen and are not trained on refusal data.Source: §B.3.4 · pp. 41–42
• The steering evaluation reports best-by-val and best-target transfer settings on JailbreakBench, AdvBench, and HarmBench, with the transferred coefficient re-selected on refusal_val only for best-by-val.The target layer and direction are shared across these two settings; test evaluation follows validation selection.Source: §B.3.4 · p. 42
• The refusal-steering appendix reports aggregated mean test ASR across 42 pairs against no-steering and native-adapter reference levels, plus per-pair ASR tables for three benchmarks and two coefficient-selection settings.The supplied passages identify the reported comparisons and table scope but do not provide values or establish a winning method.Source: §B.3.4 · p. 42, Figure 14 · p. 43

B.4.1 Additional Results – Per-Pair Performance Difference (MATCHA −Relative Depth)

• Additional figures report per-pair differences between MATCHA and Relative depth for Recall@1, Jaccard@5, and Token-Jaccard@5.The supplied passages identify the metrics and comparison but do not state the direction or magnitude of the differences.Source: §B.4.1 · p. 53, Figure 15 · p. 53, Figure 16 · p. 53

C Expressivity of the layer map

• The expressivity results state that one-to-one layer maps can be approximated arbitrarily closely at attention width d = 2, while any soft correspondence is exactly realizable when d ≥ min(Lt, Ls).Soft correspondences are row-stochastic maps with strictly positive entries, so each target layer can draw on multiple source layers.Source: C Expressivity of the layer map · p. 54
• The proof uses that αℓ is a probability vector, then multiplies the ratio bound αℓ,j ≤ e^−mαℓ,p by αℓ,p > 0 and sums.Source: C Expressivity of the layer map · p. 55
• Setting the right-hand side of (21) to ε and solving for m yields a sufficient margin for ε-accuracy.Source: C Expressivity of the layer map · p. 55
• Assumption 1 requires each leading key column’s diagonal score to exceed every other score by γ, uniformly over X; off-diagonal magnitudes and remaining columns are unconstrained.Entry (j, r) of XWk is the score assigned by detector r to source layer j.Source: Figure 18 · p. 56

C.0.1 One-to-one correspondences are contained in our family

• Assumption 1 requires a projection that scores each source layer above the others in its stack, with a margin that does not degrade across inputs.Source: §C.0.1 · p. 56
• If layer activations do not reveal their own identities, learned per-layer embeddings can supply that information to the keys.Source: §C.0.1 · p. 56
• Under Assumption 1, any one-to-one layer assignment can be approximated uniformly over source stacks to arbitrary accuracy.Increasing the logit scale tightens the approximation; the proposition requires linearly decodable layer identity with a margin uniform across inputs.Source: §C.0.1 · p. 56

C.0.2 Arbitrary soft correspondences

• Assumption 2 requires each decoding direction to return exactly one for its matching layer and zero for every other layer, uniformly across inputs.Source: §C.0.2 · p. 59
• Under Assumption 2 and d ≥ min(Lt, Ls), every strictly positive row-stochastic soft correspondence is realized exactly for all inputs.Unlike the one-to-one result, this requires exact scores for layer identity, not merely their ordering.Source: §C.0.2 · p. 59
• The soft-map construction produces the desired logits independently of the source stack, after which softmax returns the target correspondence.Source: §C.0.2 · pp. 59, 61