Abundance-Aware Set Transformer

• The paper introduces an abundance-aware extension to Set Transformers that integrates element multiplicity into attention, achieving significant efficiency and representational gains.
• It develops multiset-enhanced attention and abundance-weighted aggregation strategies that maintain permutation-invariance and promise universal approximation for abundance-sensitive functions.
• Empirical results reveal marked improvements in persistence diagram learning, statistical distance estimation, and microbiome embedding, with near-perfect performance metrics and reduced computational complexity.

An Abundance-Aware Set Transformer (AA-ST) is a neural architecture that extends the Set Transformer to explicitly handle multisets—collections of elements where each unique element is associated with a non-negative integer multiplicity or abundance. This approach is motivated by domains where the abundance or count of individual elements carries meaningful structural or semantic information, such as persistence diagram learning, statistical distance estimation, and microbiome sample embedding. AA-ST maintains permutation-invariance, augments the attention mechanism to leverage multiplicity, and achieves significant computational and representational advantages over traditional set-based or naive multiset approaches.

1. Formal Definition of Abundance-Aware Input Models

Let $X = [x_1, \ldots, x_n] \in \mathbb{R}^{n \times d}$ represent $n$ distinct $d$-dimensional elements, with abundance vector $M_X = [m_1, \ldots, m_n]^T \in \mathbb{Z}_+^n$, where $m_i$ denotes the multiplicity of $x_i$. The tuple $(X, M_X)$ thus defines a general multiset. When all $m_i \equiv 1$, this formulation reduces to a conventional set. For consistency with the principles of representation learning on sets, any function $f((X, M_X))$ must remain invariant under simultaneous permutation of rows in $X$ and entries in $M_X$ (Wang et al., 2024, Selby et al., 2022).

2. Core Architectural Elements and Attention Adaptations

Abundance-aware Set Transformers generalize set attention mechanisms to operate natively on multisets by two main strategies: (A) explicit incorporation of abundance in the attention mechanism and (B) abundance-aware aggregation at the pooling stage. Multiple implementations exist, all maintaining permutation-equivalence and strict mathematical consistency.

2.1 Multiset-Enhanced Attention

In the Multiset Transformer (Wang et al., 2024), the attention mechanism is augmented by introducing a multiplicity bias matrix. For queries $Q \in \mathbb{R}^{n \times d}$, keys/values $X \in \mathbb{R}^{m \times d}$, and multiplicities $M_Q, M_X$: $$B = \frac{(M_Q - 1)(M_X - 1)^T}{\lVert (M_Q - 1)(M_X - 1)^T \rVert_F + \epsilon}$$

$$A(Q, X) = (\operatorname{softmax}(Q X^T/\sqrt{d}) + \alpha B) X$$

where $\alpha$ is a learnable scalar and $\epsilon$ ensures numerical stability. This construction biases attention toward elements of high multiplicity and enables abundance information to propagate through attention layers.

2.2 Abundance Encodings in Embeddings

Abundance encoding can be realized as (i) a scalar scaling of feature embeddings $(1+\alpha c_i) h^0_i$, or (ii) a vector augmentation $h^0_i + u(c_i)$, where $u(\cdot)$ is a learned function (e.g., a small MLP) applied to $c_i$ (Selby et al., 2022). This embedding is carried through the entire model, facilitating abundance-sensitive representation learning.

2.3 Abundance-Weighted Aggregation

For the final invariance step, a weighted pooling mechanism is used: $z = \sum_{i=1}^n \alpha_i o_i$ where $o_i$ are the output embeddings and $\alpha_i = c_i / \sum_j c_j$ reflects normalized abundances (Yoo et al., 14 Aug 2025). Alternatively, input vectors can be replicated $c_i$ times ("replication-based weighting"), effectively causing high-abundance elements to dominate the self-attention computation without any architectural change.

3. Permutation Properties and Universal Approximation

Permutation-invariance and -equivariance are rigorously maintained in all operations. Let $P$ be a permutation matrix; for any multiset input $(X, M_X)$, the architecture ensures:

• Equivariant layers: permuting $(X, M_X)$ yields permuted hidden representations.
• Invariant pooling: the final output is invariant to such permutations (Wang et al., 2024).

The abundance-aware Multi-Set Transformer is a universal approximator for any continuous, abundance-sensitive, partially permutation-invariant/equivariant function on multisets, as formalized by Theorem 3.1 in (Selby et al., 2022). This expressive capacity stems from the ability of abundance-aware attention and feed-forward layers to approximate any such function through quantization and contextual encoding of both features and multiplicities.

4. Computational Complexity and Pool-Decomposition

Relative to naive set-transformer approaches that rely on explicit instance replication (thus incurring $\mathcal{O}(n^2 m^2)$ complexity for maximum multiplicity $m$), abundance-aware designs achieve substantial efficiency. The Multiset Transformer computes attention and pooling over the $n$ unique elements directly, reducing both time and space complexity to $\mathcal{O}(n^2)$ or $\mathcal{O}(n q)$ (with $q$ inducing points for approximate attention) (Wang et al., 2024).

Standard pool-decomposition for sets, $$f(X) = \rho(\operatorname{pool}_{i=1..n} \phi(x_i))$$, is enriched in abundance-aware models to include explicit multiplicity alignment throughout equivariant and invariant layers. This propagates abundance information and supports full leverage of multiset structure.

5. Preprocessing and Practical Recipes

DBSCAN or similar clustering algorithms can be employed to merge closely situated points and aggregate their abundances, yielding $(X', M')$ with $k \ll n$ unique elements (Wang et al., 2024). This preprocessing step can effect up to 99% reduction in effective sequence length for Transformer input, with negligible loss in task accuracy in empirical studies.

Abundance-aware Set Transformer architectures may employ various hyperparameter choices:

• Embedding dimension $d_{\rm model}$ in [128, 768], number of attention heads in [4, 12], number of induced points $m = 16$ (Selby et al., 2022, Yoo et al., 14 Aug 2025).
• Optimizers (Adam), layer normalization, dropout regularization, and batch size are set according to available GPU memory and task-specific validation performance.
• Both soft abundance-weighting and replication-based schemes are supported, with abundance-exponent ablation revealing that linear weighting ($p=1$) yields optimal results for selected microbiome tasks (Yoo et al., 14 Aug 2025).

6. Empirical Results and Applications

Abundance-aware Set Transformers have been validated in several domains:

• Persistence diagram learning: On synthetic highest-frequency-class tasks, incorporating multiplicities dramatically improves accuracy (16–56% to 41–100%; near-perfect for 2–3 classes). In TDA benchmarks (MUTAG, PROTEIN, COLLAB), the Multiset Transformer outperforms PersLay, and clustering achieves nearly same accuracy with up to 99% input reduction (Wang et al., 2024).
• Statistical distance estimation: Abundance-aware models provide finer approximation of KL divergence and mutual information compared to previous approaches (Selby et al., 2022).
• Microbiome representation: The AA-ST yields consistent and sometimes perfect classification performance compared to baseline pooling and unweighted Set Transformers in diverse phenotype and environmental prediction tasks (macro-F1 up to 1.000 on co-occurrence tasks), with explicit abundance-weighting preserving high-abundance taxa signals and allowing for complex co-occurrence modeling (Yoo et al., 14 Aug 2025).

A summary of comparative performance is provided:

Application Domain	Method	Accuracy/F1 (best)	Key Efficiency Gain
Persistence Diagrams	Multiset Transformer	41–100%	$\leq \mathcal{O}(n^2)$
Microbiome Embedding	AA-ST	1.0 / 1.0 (F1)	No core arch. changes
KL/MI Estimation	MS-Transformer	Superior to baselines	Universal approximation

7. Significance and Theoretical Implications

Explicit abundance awareness in set-based attention models bridges a critical gap in processing multisets, enabling rigorous permutation-invariant and -equivariant learning with strong theoretical guarantees. The approach supports efficient computation and universal function approximation for abundance-sensitive tasks. The methods demonstrated in (Wang et al., 2024, Selby et al., 2022), and (Yoo et al., 14 Aug 2025) establish AA-ST as a canonical architecture for domains requiring multiset modeling, including topological data analysis, statistical learning, and biological data embedding. A plausible implication is that as the field advances, further development of abundance-sensitive attention mechanisms, efficient pooling, and data-specific preprocessing pipelines can yield new capabilities and insights in multiset-based representation learning.