Neural Collapse by Design: Learning Class Prototypes on the Hypersphere

2026-05-22 · Source: cs.LG updates on arXiv.org · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Data Science & Analytics · Depth: Expert, extended

Summary

The paper "Neural Collapse by Design" introduces NTCE and NONL, two normalized loss functions, and a Fixed Prototypes (FP) approach for Supervised Contrastive Learning (SCL) to achieve Neural Collapse (NC) in supervised classification. It unifies cross-entropy (CE) and SCL under a prototype-contrast framework on the unit hypersphere, addressing CE's unconstrained radial degrees of freedom and SCL's linear probing redundancy. Empirically, NTCE and NONL surpass CE accuracy, approximate NC geometry by ≥95%, and converge in under 7.5% of CE's iterations on 4/5 metrics across four benchmarks including ImageNet-1K. FP matches SCL's linear probing accuracy, eliminating hours of training. These methods yield +5.5% mean relative improvement in transfer learning, up to +8.7% under severe class imbalance, and lower mCE on ImageNet-C.

Key takeaway

For Machine Learning Engineers optimizing supervised classification models, you should consider adopting normalized loss functions like NTCE or NONL. These methods accelerate convergence to Neural Collapse, improving generalization, transfer learning, and robustness, especially under class imbalance. If using Supervised Contrastive Learning, you can replace the time-consuming linear probing phase with fixed class-mean prototypes, significantly reducing computational overhead without sacrificing accuracy.

Key insights

Supervised learning can achieve optimal Neural Collapse geometry by unifying cross-entropy and contrastive methods on the hypersphere.

Principles

• Constrain optimization to the unit hypersphere.
• Expand negative sets for contrastive objectives.
• Decouple alignment and uniformity terms.

Method

NTCE increases negative samples from K classes to M batch instances. NONL excludes same-class instances from the denominator to resolve alignment-uniformity coupling. For SCL, use class-mean embeddings as fixed prototypes.

In practice

• Implement NTCE/NONL for faster NC convergence.
• Use fixed prototypes for SCL to eliminate linear probing.
• Target moderate temperatures (e.g., 0.1-0.2) for normalized losses.

Topics

• Neural Collapse
• Supervised Contrastive Learning
• Normalized Loss Functions
• Representation Learning
• Hypersphere Optimization
• Transfer Learning

Code references

• pakoromilas/nc_by_design

Best for: AI Scientist, Machine Learning Engineer, Research Scientist

Related on AIssential

• Beyond Neural Collapse: Task-Intrinsic Geometry Governs Neural Representations in Modular Arithmetic — Neural networks on modular arithmetic tasks form task-intrinsic cyclic geometries, deviating from neural collapse due to a structured trade-off.
• Domain-Aware Hierarchical Contrastive Learning for Semi-Supervised Generalization Fault Diagnosis — DAHCL improves semi-supervised fault diagnosis by correcting pseudo-label bias and leveraging uncertain unlabeled data through domain-aware and hierarchical contrastive learning.
• Symmetrization of Loss Functions for Robust Training of Neural Networks in the Presence of Noisy Labels — Symmetrizing multi-class loss functions, especially cross-entropy, yields a unique convex unhinged loss robust to noisy labels.
• Fixing Unsupervised Hyperbolic Contrastive Loss [D] — Hyperbolic contrastive loss requires a higher temperature parameter than Euclidean loss due to different distance scales.
• [P] Graph Representation Learning Help — Representation collapse in JEPA-style models often stems from predictor capacity or insufficient objective diversity.
• Singularity-aware Optimization via Randomized Geometric Probing: Towards Stable Non-smooth Optimization — S-Adam stabilizes non-smooth deep learning optimization by dynamically adjusting step sizes based on local geometric instability, preventing gradient chattering.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by cs.LG updates on arXiv.org.