SlowGuess/NeurIPS_2023_OCR huggingface dataset:
More complex architectures. Figure 10 shows training behavior for a more complex model than the simplified one considered in Section 5, namely where we train all parameters, replace the linear  $W_{F}$  feedforward layer by a two-layer MLP, and were (pre-)layer-normalization is added. Despite these changes, we see similar behavior for the memory recall probes (which now involve embeddings that may change over time), suggesting that the model is still identifying the same memory associations, despite the additional redundancies in parameters and modified training dynamics.
Figure 11 shows the attention maps obtained when training a multi-head version of our two-layer model, with four attention heads per layer. We see that the redundancy of multiple heads creates difficulties in identifiability: only one of the first layer heads learns the previous token behavior, while the induction behavior is shared across different heads at the second layer. This illustrates the challenges of interpretability in the presence of redundant models, which then require additional work to identify which of the layers and heads are performing a given behavior, e.g., through interventions and causal mediation analysis [35, 54].