Source-linked AI summary
Janossy Pooling: Learning Deep Permutation-Invariant Functions for Variable-Size Inputs
Ryan L. Murphy, Balasubramaniam Srinivasan, Vinayak Rao, Bruno Ribeiro
TL;DR
Permutation-invariant pooling needs to be flexible, learnable, and compatible with variable-size inputs, but naive Janossy pooling can be computationally prohibitive. The paper represents invariant functions by averaging a permutation-sensitive function over input reorderings and develops tractability strategies, reporting improved performance over current state-of-the-art methods. Its framework unifies existing approaches and supports extensions, while permutation sampling optimizes a surrogate objective rather than the original one.
Problem
Pooling operators should handle variable-size inputs while flexibly learning permutation-invariant functions beyond predefined heuristics.
Method
Janossy pooling averages a learnable permutation-sensitive function over input permutations, with tractability strategies based on k-ary interactions and random permutations.
Results
The experiments demonstrate improved performance over current state-of-the-art methods.
Takeaways & Limitations
Janossy pooling provides a framework that unifies and extends existing permutation-invariant methods, including DeepSets as a special case.
Takeaways & Limitations
Naive Janossy pooling is computationally intractable because prediction and gradient computation require summing over all permutations.
Abstract
from arXiv · showhide
We consider a simple and overarching representation for permutation-invariant functions of sequences (or multiset functions). Our approach, which we call Janossy pooling, expresses a permutation-invariant function as the average of a permutation-sensitive function applied to all reorderings of the input sequence. This allows us to leverage the rich and mature literature on permutation-sensitive functions to construct novel and flexible permutation-invariant functions. If carried out naively, Janossy pooling can be computationally prohibitive. To allow computational tractability, we consider three kinds of approximations: canonical orderings of sequences, functions with $k$-order interactions, and stochastic optimization algorithms with random permutations. Our framework unifies a variety of existing work in the literature, and suggests possible modeling and algorithmic extensions. We explore a few in our experiments, which demonstrate improved performance over current state-of-the-art methods.
1 INTRODUCTION
The paper seeks flexible, learnable pooling operators that remain invariant to input ordering and accept variable-size inputs. Janossy pooling averages a permutation-sensitive function over reordered embeddings, while tractability strategies trade computational cost against interaction capacity.
- Motivation: Pooling must be invariant to input ordering and support variable-size inputs, especially for graph neighborhoods with widely varying numbers of neighbors.Existing operators are largely predefined heuristics such as max, min, sum, and average.
- Motivation: The general architecture is a universal approximator, but learning its second function can be difficult when useful intermediate representations require higher-order interactions.For the range function y(x) = max_i,j≤|x| |x_i − x_j|, the summation can mix information that the second function must undo.
- Contributions: The authors present Janossy pooling as a framework that unifies and extends existing permutation-invariant methods, including DeepSets as a special case.They also connect Janossy pooling to finite de Finetti exchangeability.
- Janossy pooling: Janossy pooling constructs permutation-invariant outputs by applying a learnable permutation-sensitive function to every permutation of an input embedding and summing the results before a second function.The permutation-sensitive function can include feedforward or recurrent neural networks.
- Tractability: Janossy pooling exposes a choice between functions depending on the entire reordered sequence and functions depending only on their first k arguments.Restricting dependence to the first k arguments allows up to k-ary dependencies while improving tractability.
- Tractability: The framework includes permutation sampling during training, using sequence or feedforward models and the π-SGD algorithm to learn permutation-invariant networks.The paper analyzes why this procedure is theoretically justified and relates it to randomized model ensembles.
2 JANOSSY POOLING ⇀,
Janossy pooling constructs permutation-invariant functions by averaging a permutation-sensitive function over input reorderings, then offers tractable approximations that trade computational cost against flexibility.
- Janossy pooling: Janossy pooling averages a permutation-sensitive function over all permutations of a variable-size input sequence to produce a permutation-invariant function.The resulting representation can use feedforward, recurrent, or convolutional permutation-sensitive functions.
- Janossy pooling: A second function ρ can map the pooled representation to the final output and recover capacity lost by approximating or restricting the pooling function.Theoretically, ρ adds no representational power for sufficiently flexible exact pooling, but it can improve learnability and compensate for tractable restrictions.
- Tractability through canonical input orderings: Canonical ordering removes the permutation sum by sorting inputs before applying a permutation-sensitive model, but performance depends on whether that order is task-relevant.A mismatched ordering can lead to poor prediction accuracy, while learning an ordering requires searching over |h|! permutations.
- Tractability through k-ary dependencies: k-ary Janossy pooling sums over |h|!/(|h|−k)! terms, reducing computation when k<|h| while limiting the order of interactions the model can express.The hyperparameter k balances computational savings against the capacity to model higher-order dependencies; DeepSets is the k=1 case.
- Tractability through permutation sampling: Permutation sampling optimizes a surrogate objective whose optima can differ from the original objective, although the surrogate is an upper bound under convexity and identity-ρ conditions.If the function class includes permutation-invariant functions, global minima of the surrogate include those of the original objective.
- Tractability through permutation sampling: π-SGD estimates Janossy pooling by sampling random permutations during training, and its optimization enjoys almost sure convergence under conditions analogous to SGD.Using more sampled permutations reduces estimator variance, while sampling all |h|! permutations recovers the exact algorithm.
3 EXPERIMENTS
The experiments evaluate k-ary Janossy pooling and π-SGD on permutation-invariant sequence and graph tasks. Results show that higher-order interactions and permutation sampling can improve performance, with task and upper-network complexity shaping which approximation works best.
- Experimental setup: The experiments compare unary, k-ary, and full Janossy pooling, using linear or MLP upper functions and π-SGD with LSTM or GRU sequence models.The sequence models use sampled permutations at inference, while k-ary models restrict interactions to orders k = 1, 2, or 3.
- Arithmetic tasks: π-SGD models typically achieve top or statistically comparable performance across tasks, and provide significant gains over k ∈ {1, 2, 3} on the variance task.The variance task benefits from modeling full dependencies despite π-SGD optimizing an approximate loss.
- Arithmetic tasks: For MLP upper functions, k = 1 performs consistently better than k ∈ {2, 3}, while greater pooling complexity works better with a linear upper function.The authors attribute the poorer k ∈ {2, 3} MLP results to optimization issues and report that these models require more epochs.
- Arithmetic tasks: Higher-order pooling improves accuracy on range, unique sum, unique count, and variance tasks when the upper function is linear.The sum task is naturally unary, whereas the other tasks require exploiting high-order relationships within the sequence.
- Graph tasks: Increasing sampled permutations at test time reduces variance and can improve accuracy, with the graph experiments finding diminishing returns after seven samples.On PPI, increasing from one to three samples improves accuracy, while seven versus one is significant with p < 10^-3 over 12 replicates.
4 RELATED WORK
The paper places Janossy pooling within three approximation strategies for permutation-invariant modeling: canonical orderings, k-ary dependencies, and permutation sampling. It also connects the framework to exchangeability and prior invariant or relational architectures.
- Approximation approaches: Existing approaches to the intractable Janossy layer are organized as canonical orderings, k-ary dependencies, and permutation sampling.The framework is also discussed in relation to invariant models and probabilistic interpretations.
- Canonical orderings: Canonical ordering achieves invariance by mapping each input permutation to an ordered sequence, but its usefulness depends on whether that ordering is relevant to the task.Learned canonical orderings require searching among |h|! permutations, creating a difficult discrete optimization problem.
- k-ary pooling: DeepSets is unary Janossy pooling with k = 1, while related work extends invariant pooling with inductive biases or higher-order relationships.Examples include monotonicity constraints and pairwise relational reasoning.
- Permutation sampling: Permutation sampling provides a stochastic optimization procedure for Janossy pooling but learns only an approximate solution to the original permutation-invariant function.Earlier uses of random permutations were heuristic, whereas π-SGD supplies a framework for the procedure.
- Probabilistic interpretation: Janossy pooling is connected to finite exchangeability, where distributions can involve dependent random variables, unlike the conditionally independent mixture structure associated with infinite exchangeability.Higher-order pooling is designed to exploit dependencies arising under finitely exchangeable distributions.
5 CONCLUSIONS
The paper concludes that Janossy pooling unifies permutation-invariant modeling approaches and focuses on k-ary interactions and random permutations as tractable approximations. It identifies open questions about approximation behavior, optimization, and broader applicability.
- Conclusions: Janossy pooling unifies existing approaches to permutation invariance and supports methodological and theoretical extensions.The paper focuses on k-ary interactions and random permutations.
- Conclusions: Exact Janossy pooling is computationally intractable because prediction and gradient computation require summing over all input permutations.The proposed approximations create trade-offs between learnability and computational cost.
- Conclusions: Adding ρ can recover lost model capacity for tractable k-ary families, but it can reduce tractability and identifiability.Restrictions such as convexity or Lipschitz continuity may provide finer control of this trade-off.
- Conclusions: π-SGD produced the strongest empirical performance, but the relationship between its approximate loss J and the original Janossy loss L remains unresolved.The authors call for work identifying suitable problems and determining when convergence criteria are satisfied.
- Conclusions: Future work should examine canonical ordering, random-permutation optimization, and applications to harder graph tasks and non-Poisson point processes.The paper also calls for a wider range of applications.
A PROOFS OF RESULTS
The proofs establish that increasing Janossy interaction order strictly expands expressivity and that k-order pooling reduces the permutation sum to a tractable subset. They also show π-SGD converges almost surely under stated stochastic-approximation conditions.
- Computational reduction: k-order Janossy pooling therefore sums over |h|!/(|h|−k)! terms instead of all |h|! permutations.This saves computation whenever k < |h|.
- Computational reduction: Grouping permutations that agree on their first k elements yields |h|!/(|h|−k)! equivalence classes.Each class contains (|h|−k)! permutations obtained by reordering the remaining elements.
- Expressivity: k-ary Janossy pooling strictly contains (k−1)-ary pooling when H has more than one element.Thus every function representable with (k−1)-ary dependencies remains representable, but the converse does not hold.
- Expressivity: The proof characterizes equality between successive orders through a quotient that must equal one for every sequence.That condition is generally impossible unless H is a singleton, which is excluded.
- π-SGD convergence: π-SGD converges almost surely to the optimal parameters under conditions analogous to those required for SGD.The proof uses unbiased sampled gradients, independent permutations, and a supermartingale convergence argument.
B.1 RESULTS
The experiments evaluate Janossy pooling approximations across sequence and graph tasks using accuracy, RMSE, and mean absolute error. Results indicate benefits from higher-order interactions, recurrent permutation-sensitive functions, and additional sampled permutations at inference.
- Sequence tasks: Mean absolute error decreases as k increases and as more permutations are sampled at test time.For example, Janossy-20inf-LSTM outperforms Janossy-1inf-LSTM on the reported tasks.
- Sequence tasks: RNN-based f functions trained with π-SGD are especially effective on the variance task, where sequence dependencies matter.The reported results also associate more inference-time permutation samples with reduced variance.
- Sequence tasks: Table 3 reports Accuracy for sequence arithmetic tasks and RMSE for the variance task across Janossy approximations.The method column identifies the strategy for handling the permutation sum, while Infr sample records sampled permutations at test time for π-SGD methods.
- Sequence tasks: Table 4 reports Mean Absolute Error for Janossy pooling approximations under distinct tasks.Its method and Inf sample columns identify the tractability strategy and the number of sampled permutations used to estimate equation 11.
- Sequence tasks: Table 5 reports Mean Absolute Error specifically for the variance task under different Janossy pooling approximations.The table distinguishes tractability strategies and inference-time permutation samples.
- Graph tasks: Figure 2 plots graph-task performance as a function of the number of inference-time permutations.The figure concerns the PPI task.
B.2 IMPLEMENTATION AND EXPERIMENT DETAILS
The experiments implement sequence and graph versions of Janossy pooling in PyTorch, varying interaction order, sequence models, sampled permutations, and graph-neighborhood aggregation settings. Dataset summaries, parameter controls, optimization procedures, and inference details define the evaluation setup.
- Sequence tasks: For sequence k-ary models with k ∈ {2, 3}, the input sequence is sorted beforehand to reduce the combinations summed over.The embedding dimension is adjusted as part of the parameter-control design.
- Sequence tasks: The sequence implementation uses MLP or RNN functions f and linear or one-hidden-layer networks ρ.The reported architectures include tanh-activated MLPs, LSTMs with 50 neurons, and GRUs with 80 hidden neurons.
- Sequence tasks: The main k-ary results keep parameter counts consistent across k ∈ {1, 2, 3} by adjusting the embedding output dimension.Additional experiments relax this restriction and report the resulting parameter counts separately.
- Graph-based tasks: Graph experiments use PyTorch implementations following GraphSAGE, with a separately implemented LSTM aggregator.Each convolution samples k1 or k2 neighbors, combines vertex and neighbor representations, and repeats the process with new weights.
- Graph-based tasks: PPI contains multiple graphs, with 121 binary targets; training nodes come from 20 graphs and testing nodes from two unseen graphs.The dataset annotations distinguish average-graph characteristics from per-target properties.
- Graph-based tasks: The graph-task parameter count is independent of k1 and k2 by design.Variation arises from dataset-dependent input and output feature dimensions.
- Optimization: Optimization uses Adam, with dataset-specific training routines and learning rates for the sequence and graph experiments.For graph experiments, Cora and Pubmed use sampled minibatches, whereas PPI uses 10 full epochs.
- Evaluation procedure: At test time, the classifier performs 20 shuffled forward passes, averages the predicted softmax probabilities, and selects the largest average.This implements permutation sampling during inference.
C LATEXFOR JANOSSY FUNCTION MARKERS
This section provides LaTeX macro definitions for typesetting Janossy function markers and related arrow notation. The macros address environments where the notation caused compilation difficulties.
- LaTeX macros: The implementation defines a LaTeX macro for the Janossy function notation.The document instructs users to include the amsmath package.
- LaTeX macros: Additional macros define arrow variants used with the Janossy notation.The commands include dbar, harrow, and harrowStable definitions.
- LaTeX macros: The stable arrow macro is introduced because the arrow notation caused difficulties in environments such as figures.The section records the workaround through a separate LaTeX command.