Matrix Zonotopic Attention: A Context-Adaptive Value Projection for Set Transformers

2026-08-05

AI Summary

This work addresses the limitations of conventional multi-head attention in permutation-invariant set tasks, where fixed value projections hinder the modeling of complex input-dependent target mappings. To overcome this, the authors propose a context-adaptive attention mechanism that dynamically adjusts value projections using a family of matrix zonotopes—defined as a center matrix plus a weighted sum of generator matrices gated by the input—thereby enhancing representational capacity while preserving permutation equivariance. They introduce a Transform Degrees of Freedom (TDOF) metric and theoretically demonstrate that a single layer of the proposed mechanism can efficiently represent targets with high TDOF, circumventing the need for deep stacking inherent in traditional approaches. Empirical results show significant performance gains over standard attention on high-rank sparse combinatorial set prediction tasks, while matching its performance on aggregate statistical tasks.

Abstract

Multi-head attention combines an input-dependent softmax routing with an input-independent linear value projection, so the per-sample operator mapping aggregated values to outputs is the same for every input set. We study the consequences of this asymmetry for permutation-invariant set targets. We introduce the Transformation Degrees of Freedom (TDOF) of a target operator, a complexity measure counting the input-dependent directions an exact representation requires, and present a depth-separation analysis showing that context-rigid attention needs depth proportional to the target's TDOF, whereas a single layer with a context-adaptive value family can represent the same target. Building on this analysis, we propose Matrix Zonotopic Attention (MZAttn), which replaces the fixed value projection with a context-adaptive matrix-zonotope family: a centre matrix plus a sum of generator matrices weighted by input-dependent gates. The construction reduces to standard multi-head attention at initialisation, preserves permutation equivariance, and admits a data-driven reachability interpretation. Experiments on a range of set-prediction tasks are consistent with the TDOF prediction that the architectural advantage is selective: it appears on targets that depend on the input set in a high-rank, sparsely combinatorial way, and is small on aggregate-statistic targets where parameter-matched standard attention is already competitive.

Problem

• set transformers
• multi-head attention
• context-adaptive value projection
• permutation invariance
• representation expressivity

Innovation

• Matrix Zonotopic Attention
• context-adaptive value projection
• Transformation Degrees of Freedom
• set transformers
• permutation equivariance