CakeStack: a Mixture of Experts mechanism of Aggregated Attention Heads

In recent years, the use of the Mixture of Experts mechanism in LLM has spread significantly~\cite{Cai_2025}, due to its ability to store large amounts of knowledge while activating only a fraction of parameters at inference time, supposedly the ones strictly needed to perform the task requested by the user. This work experiments a novel architecture employing an attention-head-level mixture of Attention (MoA~\cite{zhang2022mixtureattentionheadsselecting}) where a subset of attentive heads, denominated "core", is always kept active during the training and inference, while the other heads are replicated and organized in experts, denominated "slices". The gating function has been designed to expose each slice to a single domain: moving away from the trend that selects different experts for each token, our gating function selects a single set of experts for the whole input, under the assumption that the domain of a single prompt is really unlikely to change across tokens. The goal of this work is to obtain a model that, with respect to an architecturally similar baseline with the same number of layers, trained on the same data, achieves better performances, is more interpretable, and has comparable computational complexity, despite having more parameters.