Idea?

A Shared rotation matrix $\mathbf{R}\in {\mathbb{R}}^{d\times d}$ learned for a single concept which maps hidden states into a feature space where the causal variable is linearly separable.

Position specific Gates $G_p\in [0,1{]}^d$ for each token position $p$.

A Mask $\mathbf{M}\in [0,1{]}^d$ that identifies which dimensions encode the causal feature. At each position the intervention mask is:

$$W_P=\mathbf{M}\circ G_p$$ $${\tilde{\mathbf{h}}}^{\ast }=(1-\sum _P{W}_p)\circ {\tilde{h}}_{base}+\sum _p{W}_p\circ {\tilde{\mathbf{h}}}_{source}$$

Loss for Gates:

$$\mathcal{L}=\sum _{k=1}^d(\sum _{p\in P}(G_p{)}_k-1{)}^2$$

Should make each dimension of the learned concept subspace be “claimed” by only one token position at any given time so the model cannot active all positions at once.