PaperlayerView exact revision on arXiv

Source-linked AI summary

The Polytopal Neural Network

A. Emilie J. Wedenborg, Anders V. Nørskov, Teresa Dorszewski, Kristoffer Wickstrøm, Morten Mørup

arXiv:2610.12004v1cs.LGcs.AI

TL;DR

Deep neural networks develop structured representations, but existing methods offer limited insight into how information is transformed across layers. PNNs constrain representations to learned polytopes and use archetypal coordinates in subsequent computation; results show near-matched classification accuracy and comparable or sometimes lower test MSE than standard VQ training.

  • Problem

    Existing explanations offer limited insight into how neural networks structure and transform information across layers, while latent representations remain difficult to interpret.

    Source: §1 · p. 1

  • Method

    PNNs learn an archetypal corpus end-to-end and enforce simplex constraints on latent representations, making each a convex combination of learned archetypes.

    Source: §1 · p. 2

  • Results

    Across classification tests, PNN accuracy remains close to unconstrained CNNs (94.46% vs 94.47% on CIFAR-10), while PNN-VQ achieves comparable, sometimes lower, test MSE than standard VQ training.

    Source: §4 · p. 7

  • Takeaways & Limitations

    Archetype mixture weights provide intrinsic local explanations, while shared archetypes provide global structure and faithfully describe constrained representations.

    Source: §1 · p. 2

    PNNs add per-iteration cost for corpus-based archetype construction and simplex inference, with overhead depending on batch size, dimensions, corpus size, archetype count, and inference procedure.

    Source: §4 · p. 8

Abstract

from arXiv · show

Understanding how deep neural networks process information remains a central challenge. Existing interpretability methods often compromise structural fidelity, rely on prespecified corpora, or explain models post-hoc. We propose Polytopal Neural Networks (PNNs), a framework that extracts distinct layer-wise aspects by enforcing a polytope-based structure that is used directly in subsequent information processing. We scale our approach using learned corpus representations and an amortized simplex inference procedure and highlight how the framework also gives a direct route to vector quantized (VQ) training. In PNNs, observations are explicitly described by their alignment with layer-specific aspects. Empirical results show that imposing polytopal constraints on neural network representations preserves meaningful structures in the latent space with minimal degradation in performance, favorable compressed representations when compared to VQ representations in unsupervised learning, while also providing a performant new approach to VQ deep learning training. Our findings suggest that deep networks can enforce interpretable polytope-based representations, offering a principled path toward more transparent AI systems with minimal performance compromise.

1 INTRODUCTION

Deep networks develop structured representations, but existing explanations often reveal little about how information is transformed across layers. PNNs constrain representations to learned polytopes and use their archetypal coordinates as part of the computation.

1 INTRODUCTION

  • PNNs constrain latent features to learned polytopes, expressing each representation as a convex mixture of archetypes at every constrained layer.The simplex constraints impose non-negativity and sum-to-one conditions.Source: §1 · p. 2
  • Mixture weights provide intrinsic local explanations, while shared archetypes describe global structure and faithfully represent constrained activations.These explanatory components are part of the computation producing the output, rather than estimated afterward.Source: §1 · p. 2
  • The coordinates that explain a representation are also the coordinates passed to the next layer.This integrates interpretability into information processing rather than adding it after training.Source: §1 · p. 2
  • PNNs scale through a compact corpus learned jointly with the network and updated from mini-batches, plus amortized simplex inference.Source: §1 · p. 2
  • Replacing simplex coordinates with one-hot coordinates yields a corpus-grounded VQ model trained without EMA updates, commitment loss, or a straight-through estimator.Source: §1 · p. 2

2 RELATED WORK

Prior work uses archetypes, prototypes, and other latent constraints to explain representations, but often applies archetypal structure only at a bottleneck or as a penalty. PNNs instead enforce learned polytopal structure across network depth.

2 RELATED WORK

  • Many nonlinear Archetypal Analysis methods place archetypes at one autoencoder bottleneck or impose them through a penalty, rather than as strict constraints.AAnet used deep autoencoders to map data into a latent simplex and recover extremal structure in nonlinear domains [van Dijk et al., 2019].Source: §2 · pp. 2–3
  • The related-work gap is simplex-constrained representations throughout network depth, with explanations consistent with information processing.Post-hoc probing can find stable archetypes without changing the network, while deep Archetypal Analysis has generally used a single bottleneck.Source: §2 · p. 3
  • Prototype methods emphasize typical instances or class/part associations; archetypes instead represent convex-hull extremes, capturing boundary phenomena and diverse variation.PNNs use archetypes without class assignments, applying the constraint at multiple classifier depths, in unlabeled autoencoders, and to pretrained-backbone tokens.Source: §2 · p. 3
  • Unlike a Dirichlet VAE's fixed standard-simplex projection or linear subspace projections, PNNs learn a polytope that supports interpretable latent geometry.Source: §2 · pp. 3–4

3 METHODS

PNNs project layer activations onto learned archetypes using simplex-constrained mixtures. A learned corpus and amortized inference procedure scale this framework, which also supports token-level projections and VQ variants.

3.1 Archetypal Analysis and the Polytopal Neural Network

  • At each constrained layer, row-stochastic archetypes and sample weights define a convex-polytope projection, solved as a per-sample convex quadratic program.Archetypes are convex combinations of data, and each observation is represented as a mixture of those archetypes.Source: §3.1 · p. 4, §3 · p. 4

3.2 Scaling the PNN

  • A learned corpus replaces full-dataset archetype construction, while amortized inference and mini-batch updates reduce gradient memory dependence on the full dataset.The batch sums carry gradients; complementary sums are computed without gradients, and the resulting memory scales with batch size and corpus size rather than N.Source: §3.2 · p. 4

3.3 Theoretical Properties

  • For positive-definite G, simplex inference has a unique minimizer and Φ is continuous piecewise affine; a finite ReLU network with sparsemax represents it exactly.The sparsemax representability result is for a fixed polytope; the described implementation uses softmax.Source: §3.3 · p. 5

3.4 Token-level PNN.

  • For pretrained vision backbones, one shared polytope represents spatial tokens; their mean projected representation stays in the polytope, and token logits contribute additively to prediction.Each token receives its own simplex coordinates, and the classifier acts on the mean of the token projections.Source: §3.4 · p. 5

3.5 Generalizations of PNN

  • Changing the constraints yields VQ and unconstrained projection variants; PNN-VQ uses one-hot centroid assignments and offers gradient-based training without a commitment loss.The PNN-VQ approach retains the softmax parameterization for archetypes while estimating assignments directly.Source: §3.5 · pp. 5–6
  • In PNN-VQ, centroid archetypes are averages of assigned corpus observations, while encoder and downstream layers learn through corpus features and projected representations, respectively.The centroid solution applies provided each centroid receives at least one corpus point.Source: §3.5 · p. 5

4 RESULTS AND DISCUSSION

Across classification, reconstruction, and token-level settings, PNNs generally approach unconstrained-model performance while structuring representations as archetypal mixtures. The experiments also assess scalable inference, direct PNN-VQ training, and how archetypes organize inputs and contribute to predictions.

4 RESULTS AND DISCUSSION

  • Mini-batched training with batch-wise corpus updates achieves accuracy comparable to full-data training, and periodic realignment makes amortized simplex inference a fast, accurate replacement for full QP projection.The amortized inference procedure is evaluated as a replacement for full QP projection.Source: §4 · pp. 6–7
  • PNN-VQ achieves test MSE comparable to standard VQ training with a commitment loss, and in several cases achieves lower MSE.Detailed results are reported in Table 12.Source: §4 · p. 7
  • PNN-CNN accuracy generally approaches that of the unconstrained CNN as archetype count increases, although the remaining gap varies across datasets.The largest improvements occur when moving away from the smallest archetype counts; encoder dimension remains fixed as K increases.Source: §4 · p. 7
  • With K = 25, PNN test accuracy is 89.1% on PathMNIST, 77.4% on DermaMNIST, and 98.1% on BloodMNIST, versus 90.4%, 78.6%, and 98.3% for the corresponding unconstrained CNNs.Increasing K from 10 to 25 improves mean accuracy on PathMNIST and DermaMNIST, while BloodMNIST remains at 98.1%.Source: §4 · p. 7
  • PNN-AE generally has lower reconstruction MSE than DirVAE, while unconstrained AE and PNN-PL generally have lower MSE than PNN-AE.The reconstruction cost associated with the convex constraint generally shrinks as K increases; softmax and sparsemax achieve similar MSE across MNIST, FashionMNIST, and CIFAR-10 under the same refinement protocol.Source: §4 · p. 7
  • Token-level PNNs on pretrained ConvNeXt-T retain most of the unconstrained model’s accuracy on ImageNet-100, with little accuracy variation for K ≥100 and overlapping error bars.The PNN projects spatial tokens onto a shared polytope before averaging the resulting representations for classification.Source: §4 · pp. 6–7
  • 94.46% test accuracy on CIFAR-10 with K = 10, compared with 94.47% for the unconstrained classifier.The PNN head was added to a frozen compact vision transformer.Source: §4 · p. 7
  • FashionMNIST archetypes form recognizable clothing and footwear groups, while ImageNet token coordinates expose spatial archetype weights that can be related to class scores.For the snake example, head tokens emphasizing an archetype associated with a hognose snake scale pattern influence predictions; unlike prototype networks, PNN uses each token’s coordinates and requires no prototype-specific losses.Source: §4 · pp. 7–8

5 CONCLUSION

The conclusion presents PNN as an architecture that enforces archetype-constrained latent representations. It reports semantically meaningful explanations and a connection to vector quantization, with archetypes serving as a codebook for classification and transparent reasoning.

5 CONCLUSION

  • PNNs enforce archetype-constrained latent representations and are presented as producing semantically meaningful explanations through an architecture connected to continuous quantization.The conclusion describes learned archetypes as a semantically meaningful codebook for classification and transparent reasoning.Source: §5 · pp. 8–9

A PROOFS

The appendix establishes properties of the simplex projection and its ReLU representation, then describes scalable projection updates and evaluates amortizer fidelity, archetype removal, and token-wise representations.

A.1 Proof of Lemma 1

  • For positive definite G, the simplex-constrained minimizer is unique, and its map Φ is continuous and piecewise affine.The proof partitions the input space into polyhedral regions where the minimizer is affine.Source: §A.1 · p. 12

A.2 Proof of Theorem 1

  • A finite ReLU network followed by sparsemax represents Φ exactly for every h in R^K.The proof uses that Φ already lies in the probability simplex, where sparsemax leaves points unchanged.Source: §A.2 · pp. 12–13

B.1 Quadratic Scaling of Projections via. SMO-Updates

  • The SMO projection updates pairs of coefficients through a one-dimensional quadratic, achieving O(K^2) complexity while avoiding cubic matrix inversion.Updates are outside the computational graph, so gradients propagate through archetype parameters rather than simplex coordinates.Source: §B.1 · p. 13

B.2 Fidelity of the Amortizer

  • All 150 tested configurations lie on the diagonal when exact SMO projections replace the amortizer, so amortization does not materially change the network’s computation.Figure 7 compares test MSE for the exact and amortized projections.Source: §B.2 · p. 13, Figure 7 · p. 14

C INTERPRETABILITY OF THE POLYTOPE

  • The archetype-removal tests compare removing the highest-weight archetype with random removal across K, measuring prediction changes for PNN-CNN and MSE increases for PNN-AE.The autoencoder test renormalizes the remaining coordinates before decoding.Source: C INTERPRETABILITY OF THE POLYTOPE · p. 13, Figure 8 · p. 14

D TOKEN-WISE REPRESENTATIONS

  • The token-based PNN projects each sequence token along the feature dimension onto a shared polytope, using a corpus of cached image-token features updated through column generation.The cached features are recomputed from the current backbone each epoch.Source: D TOKEN-WISE REPRESENTATIONS · p. 14

E.1 Data

The appendix specifies the data splits, model architectures, and variants used in the experiments, including alternative projection and quantization approaches.

E.1 Data

  • MNIST, FashionMNIST, CIFAR-10, and SVHN use official test splits with 10% of training data held out for validation; EuroSAT uses a 70/10/20 split.CNN splits are stratified by class, while autoencoder splits are unstratified.Source: §E.1 · p. 14

E.2 Architectures

  • The main CNN uses two PNN projection sites in its classifier, while the autoencoder uses a PNN bottleneck and a mirrored transposed-convolution decoder.The appendix also describes a token-wise PNN for ImageNet and a frozen-transformer head.Source: §E.2 · p. 15

E.3 Model Variants

  • The compared variants include an amortized PNN refined with SMO, one-hot PNN-VQ, unconstrained PNN-PL, VQ-VAE, and DirVAE.PNN-VQ uses nearest-archetype one-hot coordinates and the task or reconstruction loss only; VQ-VAE uses EMA, a straight-through estimator, and commitment loss.Source: §E.3 · pp. 15–16

E.4 Training Protocol

The appendix documents the training procedures and settings across the main grids and larger-scale experiments, and reports tabular benchmark results and measured training costs.

E.4 Training Protocol

  • Training proceeds from an unconstrained checkpoint through archetype initialization and a frozen-encoder phase, then updates the task network; amortizer refits use detached targets and the task loss does not update the amortizer.The main CNN and AE grids vary K over {3, 5, 10, 15, 20, 25} and use five seeds.Source: §E.4 · p. 16

E.5 Training beyond the main grids

  • The larger-scale experiments share the PNN layer, corpus selection, column generation, and a task-loss-independent amortizer, but differ in backbone, optimizer, coordinate map, and reported metric.ImageNet uses corpus selection from training tokens and deploys the amortizer without SMO refinement.Source: §E.5 · p. 17

E.6 Choice of K

  • K counts the archetypes spanning the latent polytope and trades compression against fidelity; the authors recommend choosing the smallest K with sufficient performance.Choosing separate K values by layer requires fitting many models, motivating future efficient search strategies.Source: §E.6 · pp. 17–18

E.7 Training costs

  • 4.7× the training time and 9.3× the memory are required with a full CIFAR-10 corpus versus a learned corpus; inference time does not depend on corpus size.Table 4 reports that training time increases with K while memory changes only slightly.Source: Figure 10 · p. 18, §E.7 · p. 18, Table 4 · p. 18

F LARGER BACKBONES

G ADDITIONAL EXPERIMENTAL RESULTS

  • Appendix G includes additional MLP classifier results, CNN classifier analyses and visualizations, and autoencoder reconstruction experiments.The tabular benchmark evaluation covers Iris, Wine, Breast Cancer, and Digits, spanning 4–64 features and 2–10 classes.Source: G ADDITIONAL EXPERIMENTAL RESULTS · pp. 18–19, §G.1.1 · p. 19

G.1.1 Tabular Dataset Performance

  • PNN-MLP approaches unconstrained test accuracy as K increases, degrades minimally at low K, and consistently outperforms PNN-VQ-MLP and PNN-PL-MLP.Wine and Breast Cancer achieve near-baseline accuracy with as few as five archetypes.Source: §G.1.1 · p. 19

G.1.2 Validation of Scalable Training Procedures

The scalable-training validation compares mini-batched updates and projection solvers with exact full-batch methods on three small datasets where exact computation is tractable.

G.1.2 Validation of Scalable Training Procedures

  • Mini-batched training achieves equivalent accuracy to full-batch methods, supporting the corpus-streaming approach.The comparison uses three small datasets where exact computation is tractable.Source: §G.1.2 · p. 19
  • SMO-based updates match exact QP solver performance while avoiding cubic-complexity matrix inversions.The evaluation compares CVX, QP, and SMO projection solvers.Source: §G.1.2 · p. 19
  • Treating S as constant in the simplified gradient computation introduces negligible bias.Source: §G.1.2 · p. 19

G.1.3 Training Phase Ablation

The training-phase ablation finds modest, consistent gains from end-to-end training and shows that Phase 2 performance varies across PNN variants and datasets. The CNN experiments extend comparisons to MedMNIST and to configurations with one or two constrained sites.

G.1.3 Training Phase Ablation

  • For standard PNN-MLP, the Polytope phase stays within 4 points of Pretrained accuracy on Breast Cancer, Iris, and Wine for K ≥5, but reaches 35.7 to 83.3% on Digits.Source: §G.1.3 · p. 19
  • PNN-PL maintains pretrained performance in Phase 2; on Wine, it preserves 97.8% accuracy for every K ≥5.The authors interpret this as evidence that learned subspaces capture essential structure in pretrained features.Source: §G.1.3 · pp. 19–20
  • Batched training with SMO achieves statistically equivalent test accuracy to exact methods while enabling scalability.Table 7 compares batched and non-batched training across optimization methods.Source: Table 7 · p. 20
  • For PNN-VQ, Breast Cancer accuracy falls to 86.0% at K = 3 from the 96.3% pretrained baseline, with performance recovering somewhat as K increases.PNN-VQ-MLP remains competitive on Breast Cancer, Iris, and Wine at larger K.Source: §G.1.3 · p. 20
  • Phase 3 end-to-end training provides modest but consistent improvements for standard PNN.The authors describe this as enabling full co-adaptation after pretraining and polytope-only optimization.Source: §G.1.3 · p. 20

G.2 CNN Experiments

  • The CNN experiments include PNN-CNN results on three MedMNIST datasets at K ∈ {10, 25} and compare one versus two constrained sites on MNIST, FashionMNIST, and CIFAR-10.The MedMNIST runs use five seeds; the constrained-site comparison holds the remaining configuration identical.Source: §G.2 · p. 20

G.3 AE Experiments

The AE experiments examine how polytope size and model choice relate to reconstruction quality, and compare amortizer outputs. The VQ section contrasts standard VQ-VAE and PNN-VQ reconstruction setups.

G.3 AE Experiments

G.3.1 Vector Quantization

  • Table 12 compares reconstruction MSE for standard VQ-VAE with EMA codebook, commitment loss, and STE against PNN-VQ with a corpus-grounded codebook, direct gradient updates, and reconstruction loss only, averaged over five seeds.Source: §G.3.1 · p. 23, Table 12 · p. 25
Loading 2610.12004v1…