Improving Set Function Approximation with Quasi-Arithmetic Neural Networks

AI summary

Overview

Research area: Machine learning / neural network architecture design, specifically permutation-invariant models for set-structured data (set function approximation).

Technical level: Advanced. The paper assumes familiarity with permutation invariance, Janossy pooling, invertible neural networks, and universal approximation arguments, though the core idea can be grasped intuitively.

Scope: The paper introduces a trainable version of the Kolmogorov (quasi-arithmetic) mean as a learnable pooling operation, embeds it in a new architecture family called quasi-arithmetic neural networks (QUANNs), proves approximation properties, and evaluates it against nine baselines on one synthetic and four real-world datasets.

What This Paper Is About

Most neural networks that process sets (like point clouds or collections of images) combine the elements using a fixed rule such as sum or max. Because that rule cannot be trained, the encoder is forced to shape its embeddings around an arbitrary choice made in advance, which limits expressivity and makes the learned embeddings less reusable. This paper replaces that fixed rule with a learnable, invertible transformation inside the classic Kolmogorov mean, giving the model a trainable notion of "central tendency" and, the authors argue, more structured and more transferable embeddings.

Key Contributions

1. The Neuralized Kolmogorov Mean (NKM): the authors propose what they describe as the first machine-learnable modification of the Kolmogorov mean, where the generating function is implemented by an invertible neural network, yielding a learnable measure of central tendency.
2. Quasi-arithmetic neural networks (QUANNs): a family of set function models that uses the NKM as a learnable pooling operation. Two instances are defined: QUANN-1 (unary, encoding individual elements) and QUANN-2 (binary, encoding element pairs via attention, in the style of SetTransformer).
3. Theoretical analysis: a universal approximation result (Theorem 5.1) showing QUANNs can uniformly approximate all permutation-invariant set functions that are uniformly continuous with respect to the Wasserstein metric, plus analysis of how well QUANNs approximate mean-, max- and sum-decomposable functions, and arguments that invertibility prevents dimensional collapse and yields structured latent spaces.
4. Empirical evaluation: comparisons against nine baselines across a synthetic aggregation suite, MNIST-sets, Omniglot-sets, ModelNet40 and QM9, plus embedding-quality and encoder-transfer experiments.

Main Findings

• Synthetic aggregation tasks: QUANN-1 outperformed the three ablated variants in 9 out of 10 experiments, with significant improvements in 7 of them (evaluated by mean squared error across ten different set aggregation functions).
• MNIST-sets with generic encoders: QUANN-1 outperformed all unary baselines across all 5 aggregation tasks; QUANN-2 achieved the best overall performance across all models evaluated.
• Embedding quality: after projecting encoder outputs with UMAP, only the embeddings from QUANNs and SetTransformer consistently produced class-wise separable structure across all tasks.
• Pre-trained encoder robustness: using a frozen MNIST classifier encoder that reached 98% test set accuracy, only QUANNs retained their performance across all tasks while baselines experienced significant performance drops, despite similar numbers of trainable parameters.
• Omniglot-sets: across all four cardinality settings (n_min = 5, n_max = 10, 15, 20 and 25), QUANNs surpassed all baselines on macro-averaged balanced accuracy.
• Omniglot transfer learning: encoders learned by QUANNs gave the best downstream image classification accuracy, with QUANN-1 at 0.597 and QUANN-2 at 0.589, versus SetTransformer at 0.480, DeepSet at 0.524, SlotAtt at 0.484, FSPool at 0.463, HPDS at 0.435, NormDeepSet at 0.393, PointNet at 0.256 and LAF at 0.156.
• ModelNet40: QUANNs were outperformed by FSPool and PointNet but not by other baselines — a result the authors attribute to those models' suitability for detecting extreme values in point clouds.
• QM9: QUANN-2 achieved the best performance across all baselines on both regression tasks (HOMO energy level and LUMO energy level).
• Approximation orders (Table 2): QUANNs achieve error O(1) for mean-decomposable functions by learning a linear ψ; O(1/w log n) for max-decomposable functions by learning an exponential-family ψ; and O(n) for sum-decomposable functions — meaning approximation error grows linearly with set cardinality. Expected error for sum-decomposable functions can still be made arbitrarily small when the encoder absorbs the expected cardinality.
• Aggregate picture: the paper's Figure 4 reports the relative number of outcomes where the model in a given row outperformed the model in a given column, with asterisks marking statistical significance (p < 0.05) and double asterisks very high significance (p < 0.01).

Methodology in Plain English

The starting point is the classical Kolmogorov mean, also called the quasi-arithmetic mean. It works like this: pick an invertible function f, apply it to every number in a collection, take the ordinary arithmetic average of the transformed values, then apply the inverse of f to get back to the original scale. Different choices of f recover familiar means — linear f gives the arithmetic mean, log gives the geometric mean, a power gives power means.

The authors' move is to stop choosing f by hand and instead let a neural network learn it, using an invertible architecture (they use RevNet in all experiments). Because f is invertible, it must be strictly monotonic, which guarantees the aggregate stays between the minimum and maximum of the inputs. This learned-aggregation step becomes the pooling operation in a standard set-function pipeline: an encoder maps each element (or element pair) to a latent vector, the neuralized mean aggregates those vectors, and an estimator network produces the final prediction.

For evaluation, the authors run a controlled synthetic experiment with ten aggregation targets and three ablations — replacing the learned ψ with the identity function, replacing it with the identity while enlarging the encoder and estimator to match QUANN-1's parameter count and depth, and removing the normalization term in the pooling equation. They then run the standard set-learning benchmarks with both from-scratch encoders and frozen pre-trained encoders, and finally check whether encoders learned inside the set models transfer to a plain image classification task.

Theoretical work proceeds by characterizing which function classes are uniformly approximable, and by deriving the gradient of the model with respect to encoder parameters to show why invertibility keeps latent dimensions from collapsing.

Why This Matters

For research, the paper reframes the pooling step in permutation-invariant architectures — long treated as a fixed design choice — as something learnable and principled. It connects neural architecture design to a well-studied mathematical object (the Kolmogorov mean) and provides approximation-theoretic guarantees alongside empirical results. It also argues that learning an invertible aggregator produces embeddings that transfer off the original set task, which speaks to a recurring concern about whether set models learn general-purpose representations.

Real-world applications suggested or implied:

• Graph neural networks: replacing conventional message-passing aggregators with a learnable generalized mean.
• Multi-modal and multi-view learning: learning the fusion rule that combines representations from several modalities or sources.
• Federated learning: aggregating model updates arriving from many clients.
• Point cloud and perception tasks: the paper's own experiments cover point clouds (ModelNet40) and molecular property prediction (QM9), both of which involve sets of coordinates or atoms.

Industry relevance: any pipeline where per-item embeddings must be combined into a single prediction — 3D perception, molecular property prediction, and cross-client model aggregation — can adopt the aggregation layer without redesigning the encoder, and the transfer-learning result suggests pre-trained backbones may need less downstream fine-tuning.

Future Directions

• Sum-decomposable functions: the authors identify limited capacity here as their major limitation and note that when the target function is additive or expansive, a simple sum may be a better choice than the NKM.
• Small-data regimes: the authors suggest that with limited data, the NKM may be unnecessarily flexible and prone to overfitting, whereas a fixed function offers better inductive bias and computational simplicity — an empirical question left open.
• Uncountable sets: the theory and experiments are restricted to countable sets. Extending to uncountable sets would mean replacing the sum in the pooling equation with a form of integration, but the authors state they cannot guarantee which theoretical benefits would carry over, and call set function approximation over uncountable domains an open question.
• Architectural choices: the authors report that alternative invertible architectures were explored for the generating function and did not lead to a qualitative change in performance, but the practical design space for the aggregation layer remains largely unprobed.

Target Audience

Machine learning researchers and graduate students working on set-structured data, permutation-invariant architectures, or learnable aggregation in neural networks. It will also interest practitioners building models for point clouds, molecular data, or multi-source fusion who are willing to work with invertible network components, and theorists interested in universal approximation results for set functions. Readers without a background in measure-theoretic approximation arguments will find the empirical sections far more accessible than the theory.

Authors’ abstract

Sets represent a fundamental abstraction across many types of data. To handle the unordered nature of set-structured data, models such as DeepSets and PointNet rely on fixed, non-learnable pooling operations (e.g., sum or max) -- a design choice that can hinder the transferability of learned embeddings and limits model expressivity. More recently, learnable aggregation functions have been proposed as more expressive alternatives. In this work, we advance this line of research by introducing the Neuralized Kolmogorov Mean (NKM) -- a novel, trainable framework for learning a generalized measure of central tendency through an invertible neural function. We further propose quasi-arithmetic neural networks (QUANNs), which incorporate the NKM as a learnable aggregation function. We provide a theoretical analysis showing that, QUANNs are universal approximators for a broad class of common set-function decompositions and, thanks to their invertible neural components, learn more structured latent representations. Empirically, QUANNs outperform state-of-the-art baselines across diverse benchmarks, while learning embeddings that transfer effectively even to tasks that do not involve sets.