Positional Discovery for Causal Abstraction Interventions

HyperDas

To localize a concept in the layer $l$ of a target model $\mathcal{M}$, a HyperDAS architecture consists of a hyper-network $\mathcal{H}$ that takes in a text specification $\mathbf{x}$ of the target concept and dynamically selects token position in the base text $\overline{\mathbf{x}}$ and counterfactual text $\hat{\mathbf{x}}$ and identifies a linear subspace $\mathbf{R}$ that mediates the target concept. Selecting tokens is a discrete operation, so the selection during training is “soft” and force discrete decisions during eval.

How does it work?

The hypernetwork receives a text instruction (e.g., the country of a city) and encodes it with:

• A Transformer’s Decoder;
• Cross Attention to the Target LM $\mathcal{M}$ Hidden States (base $\overline{\mathbf{x}}$ and counterfactual prompts $\hat{\mathbf{x}}$);

The Transformer’s Hypernet-work has $N$ layers. Beside the Multi-head Attention (MHA) and the Multi-Layer Perceptron (MLP), each decoder has two additional cross-attention modules to incorporate information from target model $\mathcal{M}$ processing the base and counterfactual prompts:
Let $\overline{\mathbf{h}}\in {\mathbb{R}}^{B\times L\times d}$ and $\hat{\mathbf{h}}\in {\mathbb{R}}^{C\times L\times d}$ be the stacks of base and counterfactual hidden states from the base and counterfactual input, where $L$ is the number of layers in $\mathcal{M}$, $d$ is the hidden dimension, and $B$ and $C$ are the sequence length of the base and counterfactual examples, respectively.
Two MHA, $\overline{\text{MHA}}$ and $\widehat{\text{MHA}}$ allow $\mathcal{H}$ to attend to $\overline{\mathbf{h}}$ and $\hat{\mathbf{h}}$.

Basically, each layer of the hyper-network $\mathcal{H}$ can attend to every layer of target model.

For the $p-$th decoder layer of the hyper-network $\mathcal{H}$ the three attention mechanisms are:

$${\mathbf{e}}_p^{\prime }=\text{MHA}(\mathbf{Q}({\mathbf{e}}_p),\mathbf{K}({\mathbf{e}}_p),\mathbf{V}({\mathbf{e}}_p))$$ $${\mathbf{e}}_p^{\prime \prime }=\overline{\text{MHA}}(\overline{\mathbf{Q}}({\mathbf{e}}_p^{\prime }),\overline{\mathbf{K}}(\overline{\mathbf{h}}),\overline{\mathbf{V}}(\overline{\mathbf{h}}))$$ $${\mathbf{e}}_{p+1}=\widehat{\text{MHA}}(\hat{\mathbf{Q}}({\mathbf{e}}_p^{\prime \prime }),\hat{\mathbf{K}}(\hat{\mathbf{h}}),\hat{\mathbf{V}}(\hat{\mathbf{h}}))$$

After the final transformer block is applied, the residual stream vector at last token position $e_E^{(N)}\in {\mathbb{R}}^d$ encodes information about the concept targeted for intervention and the target model’s base and counterfactual runs.

To now select the tokens in the base and counterfactual prompts that contain the concept encoded as ${\mathbf{e}}_E^{(N)}$ they construct an “Intervention Score” matrix $G\in {\mathbb{R}}^{B\times (C+1)}$ with elements ranging from 0 to 1, where $B$ and $C$ are the number of tokens in the base and counterfactual prompts.

Each element $G_{(b,c)}$ denotes the degree to which the $b$-th base token is aligned with the $c$-th counterfactual token for intervention. The additional column (i.e., that +1 in the dimensions of the $G$) corresponds to the score for retaining $b$-th base token without intervention.

In the paper example, they find that only the base token “nna” (the last token of the entity “Vienna”) receives an intervention score of 1 when paired with counterfactual token “Paris”. All other base tokens are not selected. The concept of “country” seem to be localized to the last token of the city entities.

To compute $G_{(b,c)}$ they use ${\overline{\mathbf{h}}}_b^{(l)}$ and ${\hat{\mathbf{h}}}_c^{(l)}$, the $l$-th layer residual stream representation of the target model at $b$-th base-token and $c$-th counterfactual token. First these are linearly combined:

$$\mathbf{g}(b,c)=F([{\overline{\mathbf{h}}}_b^{(l)};{\widehat{{\mathbf{h}}_c}}^{(l)}])$$

where $F(\cdot ):{\mathbf{R}}^{2d}\to {\mathbf{R}}^d$ is a linear projection that condensates the concatenated representation into the original dimension $d$.
For the extra column that indicates no intervention, the representation is simply the original base token representation:

\mathbf{g}(b, C+1) = \overline{\mathbf{h}}_b^{(l)} $$Then, each merged representation $\widetilde{\mathbf{h}}$ is weighted using the concept encoding $\mathbf{e}^{(N)}_E \in \mathbb{R}^d$ and $M$ "query" and "key" matrices $Q^{(i)} \in \mathbb{R}^{d\times M\times \frac{d}{M}}$ and $K^{(i)} \in \mathbb{R}^{d \times M\times \frac{d}{M}}$:

G’{b,c} = \frac{\sum^M{i=1}(\mathbf{e}^{(N)}_E) Q^{(i)}\cdot (\widetilde{\mathbf{h}}) K^{(i)}}{M\sqrt{d}}

$$Then$$

G = \text{ColumnSoftmax}(G’)

Basically $G_{b,c}$ measures **how strongly** base token $b$ should be aligned with counterfactual token $c$. This enables the automatic discovery of the right token to patch between base and source (e.g. "Vienna" $\leftarrow$ "Paris") Then, the intervened hidden state is built using the matrix $G$. For the $b$-th base token hidden state $\mathbf{\overline{h}}_b^{(l)}$ the intervention representation is:

\widetilde{\mathbf{h}}b^{(l)} = G{(b, C+1)}\overline{\mathbf{h}^{(l)}b} + \sum{c=1}^C G_{b,c}\hat{\mathbf{h}}^{(l)}_c

The counterfactual hidden states remain identical to the base hidden state when $G_{(b, C+1)}=1$. Conversely, if $G_{(b,c)}=1$ for a specific position $c$, the $b$-th base token is entirely replaced by the hidden vector for the $c$-th counterfactual token. ### Identifying a Linear Subspace that contains the concept It applies a $\text{MLP}$ to $\mathbf{e}^{(N)}_E \in \mathbb{R}$ in order to produce a new vector:

\mathbf{v} = \text{MLP}(\mathbf{e}^{(N)}_E)\in \mathbb{R}

While in DAS there is a fixed low-rank matrix with orthogonal columns $\mathbf{R}'$ representing a fixed subspace targeted for intervention, here, using a Householder transformation it changes $\mathbf{R}'$ conditional on $\mathbf{v}$ into a new matrix $\mathbf{R}$ that still has orthogonal columns, the householder transformation $\mathbf{H}$ is defined as:

\mathbf{H}= \mathbf{I} - 2\frac{\mathbf{vv}^{T}}{\mathbf{v}^T\mathbf{v}}

where $\mathbf{I}$ is the identity matrix. The matrix $\mathbf{H}$ is orthogonal and $\mathbf{R}'$ has orthogonal columns, which means that $\mathbf{R}'\mathbf{H}$ has orthogonal columns. Utilizing this property, we can dynamically select the subspace based on the intervention representation $\mathbf{e}^{(N)}_E$ by computing $\mathbf{R} = \mathbf{R}'\mathbf{H}$. After selecting token-positions and identifying a subspace, they intervene. For each base token $b$ at layer $l$ of the target model $\mathcal{M}$, there is a weigthed interchange intervention with the counterfactual hidden states $\mathbf{\widetilde{h}}^{(l)}$ and the low-rank orthogonal matrix $\mathbf{R}$:

\mathbf{H}_{b}^{(l)} \leftarrow \mathbf{\widetilde{h}}^{(l)} + \mathbf{R}^T (\mathbf{R}(\mathbf{\widetilde{h}}_b^{(l)})-\mathbf{R}(\mathbf{\overline{h}}_b^{(l)}))