PaperlayerView exact revision on arXiv

Source-linked AI summary

ProtoSeam: Lifting Classifier Training with Latent Gaussian Mixture Models

Robert Lampel, Timon Klein, Sebastian Sager

arXiv:2609.35174v1cs.LG

TL;DR

End-to-end classification couples the early network to the head's gradient path and trains the head only on observed embeddings. ProtoSeam inserts class prototypes at one interface and trains the two parts without gradients crossing it; it improves test accuracy across the evaluated datasets and architectures.

  • Problem

    End-to-end training couples N1 to N2's Jacobian and trains the head only on training embeddings, leaving its behavior in nearby representations unconstrained.

    Source: §1 · p. 1

  • Method

    ProtoSeam inserts learnable class prototypes at a semantic interface, using consensus to train N1 and sampled class-conditional inputs to train N2 without gradients crossing the seam.

    Source: §1 · p. 1, §3 · p. 5

  • Results

    Lifted training improved test accuracy over baseline and matched end-to-end controls on every dataset–architecture pair and seed, by up to 5.2 and 7.3 percentage points, respectively.

    Source: §5 · p. 9

  • Takeaways & Limitations

    The evaluated gains occur with the unmodified network deployed after prototypes and sampling are discarded, and the paper analyzes a mechanism involving selection of a locally more robust optimum.

    Source: §5 · pp. 9–10

    The Gaussian class-conditional modeling assumption does not hold for every arbitrary trained network.

    Source: §3 · p. 4

Abstract

from arXiv · show

We propose a lifted reformulation of supervised classification that improves the final accuracy of standard classifiers without changing the architecture at inference time. A network $N=N_2\circ N_1$ is split at a single semantic interface and one learnable prototype per class is inserted there. Training combines a quadratic consensus penalty that pulls $N_1(x)$ toward the prototype of its class with a classification loss of $N_2$ evaluated on samples drawn around the prototypes, whereat no gradient crosses the interface. At inference the prototypes are discarded and the unmodified network $N_2\circ N_1$ is used. Across CIFAR-10, CIFAR-100, and TinyImageNet with ResNet and vision transformer backbones, lifted training improves test accuracy by up to five percentage points over variants without lifting under a shared tuning protocol. Moreover, we provide theoretical justification of those results.

1 INTRODUCTION

ProtoSeam addresses two couplings in end-to-end training by inserting class prototypes at one semantic interface and training the two network parts without gradients crossing the seam. It improves final accuracy across six dataset–architecture pairs, and the paper provides theoretical analyses of the gain.

1 INTRODUCTION

  • End-to-end training couples the subnetworks through N2's Jacobian and trains the head only on training embeddings, leaving its behavior nearby unconstrained.The authors note that an ill-conditioned or rank-deficient Jacobian can distort the gradient reaching N1.Source: §1 · p. 1
  • Prior decoupled-training methods target the gradient path, while prototype and margin losses shape embeddings but retain end-to-end gradients.ProtoSeam places its prototypes at one semantically meaningful interface.Source: §1 · pp. 1, 3
  • The method inserts one learnable prototype per class at a semantic interface, training N1 toward class prototypes and N2 on samples drawn around them.No gradient crosses the interface; discarding the lifted variables restores the original network composition.Source: §1 · pp. 1–2
  • Up to five percentage points: lifting improved accuracy over matched end-to-end and unsplit baselines across six dataset–architecture pairs.The comparison used a shared hyperparameter-tuning protocol.Source: §1 · p. 2
  • The paper analyzes the gain with an interface risk-transfer bound, a second-order sensitivity analysis, and a solvable one-dimensional model.The analyses connect sampled classification loss to deployed risk, sampling to head sensitivity, and selection effects to a closed-form model.Source: §1 · p. 2
  • Class-conditional Gaussian structure is motivated by neural collapse, which describes tight clusters whose means form a Simplex Equiangular Tight Frame late in training.The proposed loss promotes similar geometry from initialization.Source: §1 · p. 2

2 RELATED WORK

Prior work decouples network layers, shapes class-prototype geometry, or models Gaussian embeddings, but ProtoSeam combines a single semantic interface with learned prototypes and a severed gradient path.

2 RELATED WORK

  • Multiple-shooting-inspired and layer-wise auxiliary-variable methods decouple training, but prior layer-wise approaches do not impose latent-representation structure.The multiple-shooting analogy is structural: unlike exact matching constraints at the solution, this method enforces matching with a penalty.Source: §2 · p. 3
  • Prototype and margin methods shape embeddings or classifier weights, while ProtoSeam's distinction is cutting the class-loss gradient to N1 and adding inter-class repulsion.The consensus term coincides with center loss, which jointly learns class centers and penalizes feature-to-center distance alongside softmax.Source: §2 · p. 3
  • L-GM also assumes per-class Gaussian features but trains end-to-end; ProtoSeam instead decouples N1 from N2 and learns prototypes rather than fixing Simplex-ETF geometry.Neural collapse describes class means arranged as a Simplex ETF; a related approach fixes that geometry and trains the backbone.Source: §2 · p. 3
  • Unlike normalizing-flow latent models, ProtoSeam's soft Gaussian penalty applies to differentiable architectures such as standard ResNets, and its reparameterization only samples training inputs for N2.Normalizing flows require invertible architectures, while ProtoSeam does not.Source: §2 · pp. 3–4

3 METHOD

The method models each class's N1 embeddings with a Gaussian around a learnable prototype, then combines consensus, sampled classification, and repulsion losses. Design choices address covariance stability and cost, lifting location, embedding dimension, initialization, and penalty scheduling.

3 METHOD

  • N1's class-conditional embeddings are modeled as Gaussian around class-specific means and covariances, while the objective trains N2 to separate their confidence regions and N1 to follow that model.The Gaussian relation is a modeling assumption, not a property of every arbitrary trained network; the consensus penalty makes it approximately true.Source: §3 · pp. 4–5
  • The empirical class covariances are re-estimated once per epoch and regularized because estimates can be singular or ill-conditioned and can collapse toward zero.Covariances are centered on empirical class means rather than prototypes to avoid folding consensus residuals into the covariance.Source: §3 · pp. 4–5
  • The inter-class repulsion penalty separates nearby prototypes, especially early in training; benchmarks use α = 2.Without repulsion, clusters usually remain close together and do not separate.Source: §3 · p. 5
  • N2's classification term depends on data only through class labels, making its subproblem an n-component Gaussian classification problem independent of the training-sample count.The method permits as many Monte Carlo draws per class as desired, independently of minibatch composition.Source: §3 · p. 5
  • Per-epoch empirical covariance estimation adds one training-set forward pass, roughly one-third overhead per epoch, plus O(nk^2) factor memory.The authors use this computation in their benchmarks to avoid added scaling hyperparameters.Source: §3 · p. 6

3.1 DESIGN CHOICES

  • Lifting before a ReLU avoids restricting embeddings to the non-negative orthant, where boundary prototypes cannot be symmetrically surrounded by the assumed Gaussian.The authors lift before such an activation or before the final classification layer, and use a single fully connected layer for N2.Source: §3.1 · p. 6
  • A nonsingular sample covariance requires more class samples than embedding dimensions; the relative-error bound is restrictive, though experiments perform better at higher dimensions.For CIFAR-100 with 500 samples per class, the bound gives k ≲ 5 at ε = 0.1; the authors conjecture repulsion and shrinking covariances help explain the difference.Source: §3.1 · p. 7
  • Forward Sweep Initialization performs best with only minor observed effects on final accuracy, while penalty annealing starts with smaller ρ to let clusters separate before consensus tightens.Forward Sweep Initialization sets prototypes to the initial network's empirical class means.Source: §3.1 · p. 7

4 EXPERIMENTS

ProtoSeam is evaluated against both a conventional baseline and an end-to-end unlifted control using shared training and tuning protocols. Across the tested datasets and backbones, lifting improves accuracy, with the TinyImageNet ViT-S result exceeding the baseline by more than five percentage points.

4.1 PROTOCOL

  • The common protocol compares baseline, ProtoSeam, and an end-to-end unlifted control, with the control isolating lifting from the architecture change.All variants share data splits, epoch budget, and optimizer family; learning rate and weight decay are tuned by grid search.Source: §4.1 · pp. 7–8

4.2 DATASETS AND ARCHITECTURES

  • Experiments cover ResNet and custom ViT-S backbones on CIFAR-10, CIFAR-100, and TinyImageNet.These datasets and models are described as large enough for meaningful accuracy gains yet small enough for broad hyperparameter comparisons.Source: §4.2 · p. 8
  • The CIFAR-100 ViT-S sweep varies ρmax at k = 32 and k at ρmax = 16; accuracy rises steeply before stabilizing, and the study uses ρmax = 16 and k = 32 thereafter.The selected settings were found to work best across all datasets.Source: Figure 3 · p. 8, §4.1 · p. 8

4.3 RESULTS

  • On TinyImageNet ViT-S, lifted training initially trails while N1 is noisy and prototypes are positioned, then overtakes both alternatives and maintains its lead.Figure 4 plots validation accuracy over epochs, showing the mean of three seeds and their full range.Source: §4.3 · p. 8, Figure 4 · p. 9
  • Lifted training beats both baseline and unlifted variants across the tested datasets and architectures; on TinyImageNet ViT-S, accuracy improves by more than five percentage points over baseline.Table 1 reports test accuracy across both architectures and all three datasets, with lifted means beating both comparators at every seed.Source: §4.3 · p. 8, Table 1 · p. 9

5 CONCLUSION

The paper introduces prototype-based lifting at one network interface, discarding the prototypes at inference. Across three image benchmarks and two backbones, the method improves test accuracy over baseline and unlifted controls, and the authors analyze a mechanism for the gains.

5 CONCLUSION

  • Across CIFAR-10, CIFAR-100, and TinyImageNet with ResNet-8 and ViT-S, lifting improves test accuracy over baseline and unlifted controls on every dataset–architecture pair and seed.The reported gains reach 5.2 percentage points over baseline and 7.3 percentage points over the unlifted control.Source: §5 · p. 9
  • The method splits the network at one interface, inserts a Gaussian prototype per class, and discards prototypes at inference to deploy the unmodified network.The objective trains N1 toward class prototypes and N2 on samples drawn around them, without gradient crossing the seam.Source: §5 · p. 9
  • The authors identify a mechanism in which lifting selects a different, locally more robust local optimum rather than optimizing the same objective faster.A one-dimensional model exhibits the resulting effect in closed form.Source: §5 · p. 10
  • The experiments use established public CIFAR-10, CIFAR-100, and TinyImageNet benchmarks, with no human subjects or sensitive data.The authors report no application risk beyond risks already associated with image classifiers generally.Source: ETHICS STATEMENT · p. 10

A TECHNICAL DETAILS

The technical details specify how hyperparameters were selected and how image data were prepared for training and evaluation.

A TECHNICAL DETAILS

  • Learning rate and weight decay were selected by grid search for each dataset and model, using the highest validation accuracy after 200 epochs.The grid included four learning rates and five weight decays; unlifted variants used the baseline hyperparameters.Source: A TECHNICAL DETAILS · p. 14
  • Table 2 lists the final confirmed best hyperparameters for each dataset, model, and variant combination.The selected values follow the grid-search procedure.Source: Table 2 · p. 14, A TECHNICAL DETAILS · p. 14
  • Training uses random crops, horizontal flips, RandAugment, normalization, and random erasing, while validation and test data receive normalization only.Random erasing uses probability 0.5 and area fraction 0.02–0.33; evaluation splits receive no stochastic augmentation.Source: A TECHNICAL DETAILS · p. 14

B WHY LIFTING IMPROVES ACCURACY

The theoretical account distinguishes statistical selection from optimization and explains the sampled loss through interface mismatch and head sensitivity. It also presents a one-dimensional case where the selection effect is computed in closed form.

B WHY LIFTING IMPROVES ACCURACY

  • The sampled classification term equals reconnected-network risk up to an interface-mismatch term whose components are bounded by quantities controlled by the method.The passage cites Theorems B.3 and B.4 for this account.Source: B WHY LIFTING IMPROVES ACCURACY · p. 14
  • Sampling adds a non-negative smoothing penalty that acts to second order as a curvature penalty on the head, favoring flatter loss curves around prototypes among heads with equal prototype loss.This description applies in the small-perturbation regime and is attributed to Theorem B.5.Source: B WHY LIFTING IMPROVES ACCURACY · p. 14
  • The section treats statistical selection as a possible source of gains while not ruling out easier optimization, then gives a small instance with a closed-form selection effect.The closed-form analysis appears in Section B.3.Source: B WHY LIFTING IMPROVES ACCURACY · p. 14

B.1 THE SAMPLED LOSS IS THE DEPLOYED RISK UP TO INTERFACE MISMATCH

The sampled classification loss approximates the reconnected network’s deployed risk up to interface mismatches that the lifted objective controls. Sampling also adds a curvature penalty to the head’s loss around class prototypes.

B.1 THE SAMPLED LOSS IS THE DEPLOYED RISK UP TO INTERFACE MISMATCH

  • The sampled classification loss equals deployed risk up to interface-mismatch terms bounded by quantities controlled by the lifted objective.The bound is sharper when class-conditional embeddings are close to Gaussian; concentrating embeddings around prototypes shrinks mismatch for any distributional shape.Source: §B.1 · pp. 15, 17
  • The transfer bound accounts for mean, non-Gaussianity, and covariance mismatch, tying the risk gap to consensus, distributional modeling, and covariance estimation.The covariance-estimation control treats refreshed embeddings as independent draws, an idealization because N1 is fitted to the same data.Source: §B.1 · pp. 16–17
  • The variance floor trades off deliberate over-dispersion against covariance-estimation error, and the bound favors an intermediate floor.The sampling law deliberately has more variance than the embedding law, while the floor shrinks the estimation term.Source: §B.1 · p. 17
  • The risk bound controls cross-entropy rather than accuracy directly, and its population quantities leave generalization to the gap between population and training consensus.Cross-entropy bounds misclassification probability only up to a factor of 1/log 2.Source: §B.1 · p. 17

B.2 SAMPLING IS AN EXPLICIT CURVATURE REGULARIZER

  • For the linear head with cross-entropy, sampling around prototypes adds a non-negative second-order curvature penalty to the prototype loss.The paper distinguishes this statistical selection mechanism from the possibility that the lifted objective is simply easier to optimize.Source: §B.2 · p. 17

2. Consequently

Sampling can select among heads that fit prototypes equally well by favoring lower sensitivity around them. The paper demonstrates this selection effect in a one-dimensional binary example and qualifies when the curvature interpretation applies.

2. Consequently

  • Sampling selects heads whose loss curves least around prototypes, in directions where inference embeddings may deviate from them.The second-order interpretation requires small logit perturbations, and prototype curvature is not a global sensitivity measure.Source: §2 · p. 18
  • In the idealized case of matched means and covariances, the extra regularization over end-to-end training comes from the variance floor, not the full sampling covariance.Outside that case, mean and covariance mismatches also enter the comparison.Source: §2 · p. 18
  • The variance floor preserves isotropic smoothing when within-class covariance collapses, but weight decay is also needed to prevent head saturation from making the curvature penalty vanish.Scaling up the head can drive the penalty toward zero, a regime outside the expansion’s validity.Source: §2 · pp. 18–19

B.3 AN EXACTLY SOLVABLE ONE-DIMENSIONAL INSTANCE

  • In the one-dimensional, two-class example, neighborhood sampling resolves a head parameter left undetermined by fitting the prototypes and selects the least input-sensitive solution.The example fixes the first subnetwork and concerns only the classification problem solved by the head; it does not claim conditioning alone improves generalization.Source: §B.3 · p. 19

B.3.1 SETUP

  • The example uses balanced binary labels, fixed prototypes at −1 and +1, and a nonlinear head family whose members all reproduce those prototypes exactly.The parameter is invisible to prototype-only criteria but changes the head’s behavior near the prototypes.Source: §B.3.1 · p. 19

B.3.2 THE SELECTION RESULT

Prototype-only fitting leaves the head’s local sensitivity undetermined, while sampling around prototypes selects it. In the small-noise limit, the selected head is stationary at both prototypes.

B.3.2 THE SELECTION RESULT

  • For every a ∈ R, Lpoint(a) = 0, so prototype loss does not identify the head parameter.Source: §B.3.2 · p. 20
  • For every σ > 0, Lsample(a) = σ2 + 2a (2σ2 + 3σ4) + a2(4σ2 + 39σ4 + 15σ6) is strictly convex and has a unique minimizer.Source: §B.3.2 · p. 20
  • Prototype fitting leaves the head’s sensitivity undetermined, but sampling any neighborhood fixes it; in the small-noise limit, the selected value makes the head stationary at both prototypes.Source: §B.3.2 · p. 20

B.3.3 INTERPRETATION

In the one-dimensional model, sampling acts as a sensitivity regularizer and selects a head matched to the noise level. This is a local mechanism, not a general claim of improved clean accuracy.

B.3.3 INTERPRETATION

  • The small-noise sampled solution a = −1/2 cancels the linear response to deterministic reconnection error, reducing squared loss from Θ(δ2) for the identity head to Θ(δ4).Source: §B.3.3 · p. 21
  • The limiting head a = −1/2 loses roughly seven accuracy points at σ = 0.5, so the sensitivity result does not establish improved clean accuracy.Source: §B.3.3 · p. 21
  • Sampling selects a head matched to the noise level: |a⋆| shrinks as σ grows, and the σ-matched head stays within 0.15 accuracy points of identity for σ ≤ 0.5.Source: §B.3.3 · p. 21
  • The cubic model’s conclusion is deliberately local and should not be extrapolated to large |z|, where its behavior is not representative of a practical classifier.Source: §B.3.3 · p. 21
  • The analysis shows that the changed lifted objective selects a different, locally more robust head; faster optimization of the same objective is not the mechanism established here.Source: §B.3.3 · p. 21

B.3.4 ILLUSTRATION AND NUMERICAL VERIFICATION

Numerical checks agree with the model’s closed forms across the tested noise levels. The figures illustrate the head family and the response to bounded reconnection error.

B.3.4 ILLUSTRATION AND NUMERICAL VERIFICATION

  • The verification draws n = 2 · 105 samples per noise level, repeats each configuration over 25 seeds, and evaluates losses and accuracies on a held-out sample of the same size.Source: §B.3.4 · p. 21
  • Across σ ∈ [0.02, 1], estimated optima and measured sensitivities agree with the closed forms to within one seed standard deviation, and empirical loss agrees within Monte-Carlo error.Source: §B.3.4 · p. 21, Figure 5 · p. 22
  • Figure 6 compares loss under bounded deterministic reconnection error near the prototypes, showing Θ(δ2) for the identity head and Θ(δ4) for g−1/2 at σ = 0.2.Source: Figure 6 · p. 22

Related papers

Loading 2609.35174v1…