Partial FC: Training 10 Million Identities on a Single Machine

Xiang An; Xuhan Zhu; Yang Xiao; Lan Wu; Ming Zhang; Yuan Gao; Bin Qin; Debing Zhang; Ying Fu; Jiankang Deng

Source-linked AI summary

Partial FC: Training 10 Million Identities on a Single Machine

Xiang An, Xuhan Zhu, Yang Xiao, Lan Wu, Ming Zhang, Yuan Gao, Bin Qin, Debing Zhang, Ying Fu, Jiankang Deng

arXiv:2010.05222v4Oct 11, 2020cs.CVcs.DC

TL;DR

Millions of identities make full-softmax face-recognition training costly in classifier storage, logits, computation, and communication. Partial FC preserves every positive center while sampling negatives and sharding the classifier, achieving massive-scale training with competitive accuracy and robustness. Its main scope boundary is that exact full-class model parallelism remains limited by global-batch logit memory as GPU counts increase.

Problem
Full softmax becomes impractical for millions of identities because classifier storage, logits, computation, and class-dependent communication grow with the class count.
Method
Partial FC preserves every positive class center, samples only negative centers, and combines this approximation with permanent classifier ownership across GPUs.
Results
PFC scales to 64 million classes while moderate sampling maintains competitive recognition accuracy and improves robustness under conflicts, label noise, and long-tailed data.
Takeaways & Limitations
Positive-preserving negative sampling provides a scalable foundation for accurate and robust identity classification at massive class scales.
Takeaways & Limitations
Exact class-sharded full-softmax training remains limited by logit memory because each rank processes the global batch, so adding GPUs alone does not ensure scalability.

Abstract

from arXiv · show

Training face recognition models with millions of identities is challenging because classifier storage, logit memory, and computation grow linearly with the number of classes, eventually making full softmax impractical even when the backbone itself fits comfortably in memory. We present Partial FC (PFC), a scalable approximation to large-class softmax that preserves every positive class center while activating only a sampled subset of negative centers in each mini-batch. This asymmetric treatment retains every target term while avoiding exhaustive interaction with millions of mostly uninformative negatives. Our distributed implementation partitions the classifier across GPUs and samples within each owned shard, so sampling reduces local matrix multiplication and logit storage while sharding eliminates class-gradient synchronization across workers. Together, these properties reduce GPU-resident classifier memory, computation, and class-dependent communication without feature-based hard-negative retrieval. End-to-end system benchmarks demonstrate efficient scaling to massive class spaces, including 64 million classes on a single eight-GPU machine. Across large-scale face-recognition datasets, moderate sampling rates maintain competitive recognition accuracy while substantially improving training efficiency. Our best PFC configurations achieve 97.2\% TAR on IJB-C at FAR $=10^{-4}$ and 94.0\% TAR on ICCV21-MFR at FAR $=10^{-6}$. Beyond clean training data, PFC is robust to inter-class conflicts, label noise, and long-tailed identity distributions: under 40\% label noise, PFC with conflict filtering raises ICCV21-MFR TAR from 43.9\% to 80.2\%, while on long-tailed data PFC improves TAR from 87.4\% to 92.0\%. These results establish positive-preserving negative sampling as an effective foundation for scalable, accurate, and robust identity classification.

1 Introduction

Large-scale face recognition makes full-class softmax increasingly costly because classifier storage, logits, computation, and communication grow with the identity count. Partial FC preserves every positive center while sampling negatives, and its distributed design enables scaling with competitive accuracy and robustness.

Full softmax retains a d × C classifier and evaluates interactions with all classes, creating storage, computation, and gradient-synchronization costs at millions of identities.
Partial FC preserves every positive class center while approximating only the negative denominator through sampled negative centers.
The distributed implementation combines permanent class ownership with batch-level negative sampling to reduce active classifier storage, logits, computation, and class-dependent communication without feature retrieval.
PFC is validated through scaling to 64 million classes, competitive face-verification accuracy, and robustness under conflicts, label noise, and long-tailed data.

2 Related Work

Prior face-recognition systems established identity classification and margin-based softmax, but large-class training remains constrained by classifier and logit costs. Related approaches reduce or restructure class interactions through retrieval, decomposition, or class sharding.

Face-recognition supervision expanded from roughly half a million identities in Celeb-500K to millions of identities and hundreds of millions of images in WebFace260M.
HF-Softmax retrieves active classes with a randomized hash forest, whereas Softmax Dissection reduces redundant inter-class computation through objective decomposition.
Class-sharded model parallelism assigns classifier columns to ranks and uses collective reductions to recover the exact global softmax, but classifier-gradient communication remains costly.

3 Method

The method section frames softmax classification and contrastive learning as sharing a candidate-matching structure, then motivates positive-preserving approximation for persistent class-center classifiers.

The section derives a shared candidate-matching structure for softmax classification and contrastive learning while distinguishing how their candidates are represented and updated.
It analyzes exact class-sharded model parallelism and its memory limits before introducing Partial FC and its distributed implementation.
The method focuses on classifiers with persistent class centers, where the candidate bank can be explicitly sharded and sampled.

3.1 Softmax Classification and Contrastive Learning as Candidate Matching

Softmax and contrastive objectives share a row-wise candidate structure in which positives receive target terms and negatives contribute through probabilities, but their candidate representations and update carriers differ. PFC applies positive-preserving negative sampling to persistent class centers, retaining all positives while activating only sampled negatives.

Softmax classification: A softmax classifier’s persistent weights become discriminatively learned class centers through attraction to same-class features and repulsion from other-class features.
Partial FC: PFC retains persistent class centers but activates only a subset in each iteration, with the shared active set containing every mini-batch target.
Partial FC: Equation 4 encodes the key asymmetry: all positive centers receive attractive updates, while only sampled negatives receive repulsive updates and unsampled centers remain in optimizer state.
Partial FC: Because probabilities are normalized over the sampled active set, PFC gradients are generally not unbiased full-softmax gradients; resampling instead refreshes negative coverage.
Contrastive learning: Softmax and contrastive learning share score-level structure, but PFC updates persistent centers whereas CLIP updates shared encoders through transient candidates and both directions.

3.2 Problem Formulation

Full softmax becomes impractical at million-class scale because logits and classifier state grow with the class count. Class-sharded model parallelism preserves exact normalization while distributing classifier storage and computation across ranks, but its global-batch logits create a remaining memory bottleneck.

3.2.1 Model Parallelism: Class-sharded model parallelism partitions 𝑊 across 𝑘 ranks, each holding 𝑁=𝐵/𝑘 local samples and its own classifier shard.
3.2.1 Model Parallelism: Each rank computes logits only for its resident shard, while all-reduced maxima and exponentiated sums recover the exact global softmax without gathering global logits.
3.2.1 Model Parallelism: Full softmax computes 𝐵×𝐶 logits and stores a 𝑑×𝐶 classifier, making both impractical when 𝐶 reaches millions.
3.2.1 Model Parallelism: Classifier gradients remain local, while reduce-scatter sums feature-gradient contributions and returns each rank its original local-batch gradients.
3.2.1 Model Parallelism: Every collective moves at most 𝑁𝑘×𝑑 values and is independent of 𝐶, avoiding communication of global logits and classifier parameters.
3.2.2 Memory Limits of Model Parallelism: Sharding reduces per-GPU classifier storage to 𝑀W=4𝑑𝐶/𝑘, but every rank still processes the global batch, leaving local logit memory at 𝑀logit=4𝑁𝐶.
3.2.2 Memory Limits of Model Parallelism: At 80 GPUs, the logit-to-weight ratio reaches 10 for 𝑁=64, 𝑑=512, and approximately 125K classes per GPU.
3.2.2 Memory Limits of Model Parallelism: Per-GPU FC memory rises from 1.3 GB at 8 GPUs to 5.9 GB at 80 GPUs, while the logit share increases from 40% to 87%.

3.3 Approximation Strategy

PFC preserves every positive class while sampling negatives within permanently owned classifier shards. This reduces active classifier memory, logits, computation, and communication while retaining a softmax over the sampled union.

Batch-Level Positive-Preserving Sampling: PFC forms a batch-shared active set by combining every positive class with a sampled subset of negatives.The active set contains at most rC negative classes, approximately rC when the negative pool is large.
Batch-Level Positive-Preserving Sampling: Every target class receives an update, while unsampled negatives remain unchanged and negative coverage refreshes across iterations.The sampled denominator yields a biased full-softmax gradient, while selected hard negatives receive larger gradients through their softmax probabilities.
Distributed Partial FC: Each rank permanently owns a class shard, retains all local positives, and fills its activation budget with locally sampled negatives.The active classifier is gathered logically across ranks but remains distributed rather than materialized as one global matrix.
Distributed Partial FC: Global reductions compute the exact softmax over the union of sampled shards, while feature gradients return to their source ranks and classifier gradients remain local.With r=1, the procedure becomes conventional full-class model parallelism.
System Scaling: PFC maintains higher throughput and slower memory growth than data parallelism and full-class DTensor model parallelism, remaining executable at 64M classes.The comparison holds the backbone fixed and uses r=0.1 for data parallelism and PFC, versus r=1 for DTensor model parallelism.
Matrix Multiplication and Communication Efficiency: MP-PFC reduces per-rank traffic to 19.7M elements versus 64.1M for DP-PFC and 192.5M for full model parallelism.Sampling and sharding reduce local classifier traffic from Cd to rCd/k; MP-PFC also cuts weight reads 10× relative to full model parallelism.

4 Experiments

Experiments evaluate PFC across recognition benchmarks, system profiles, sampling rates, and imperfect-data settings. PFC maintains competitive accuracy while reducing classifier computation and memory, with stronger benefits under conflicts, noise, and long-tailed identities.

Benchmark Results: PFC benchmark tables compare face-verification performance across datasets, backbones, methods, and sampling rates using standardized operating points.Table 1 uses IR50 and Table 2 uses IR100; ICCV21-MFR and IJB-C report TAR at specified FAR values, while MegaFace also reports identification accuracy.
Communication Profiling: 713 ms at 32M identities: MP-PFC spends 3.5× less compute-kernel time than DTensor and 5.2× less than DP-PFC.The gap widens monotonically with identity count in the 8×H800 profiling sweep.
Communication Profiling: 64M identities: MP-PFC completes while DP-PFC and DTensor fail; MP-PFC fails only at 128M.The implementations share the backbone, optimizer, and batch size, so the difference is attributed to classifier resident state.
Communication Profiling: 194–224 ms: MP-PFC memcpy remains nearly flat through 64M identities, while DTensor rises from 234 ms at 1M to 847 ms at 32M.DTensor materializes and redistributes sharded logits, whereas MP-PFC keeps the sampled block rank-local.
Ablation Studies on WebFace: 86.3→86.9, 91.7→91.8, and 93.9→94.0: moderate sampling rates match or exceed full-class MFR-All baselines across WebFace4M, WebFace12M, and WebFace42M.For r∈[0.2, 0.4], the same ordering holds across all four demographic groups, although masked-track results are less regular.
Ablation Studies on WebFace: Moderate sampling retains clean-data accuracy and becomes more advantageous under inter-class conflicts, label corruption, and long-tailed identities.The experiments characterize partial activation as both a computational approximation and an effective regularizer for imperfect web-scale data.

5 Conclusion

The conclusion presents Partial FC as a positive-preserving approximation that scales large-class softmax by retaining targets and sampling negatives. It reports reduced classifier costs, competitive accuracy, and improved robustness on noisy and long-tailed data.

5 Conclusion: Partial FC retains every target class while activating only a sampled subset of negative classes.Combined with permanent class ownership, this reduces classifier computation and logit memory at massive class scales.
5 Conclusion: Moderate sampling rates maintain competitive face-recognition accuracy and improve robustness under noisy and long-tailed training data.