PaperlayerView exact revision on arXiv

Source-linked AI summary

Routing the Lottery: Adaptive Subnetworks for Heterogeneous Data

Grzegorz Stefanski, Alberto Presta, Michal Byra

arXiv:2601.22141v1cs.AIcs.CVcs.LG

TL;DR

Most Lottery Ticket Hypothesis methods use one universal sparse mask, overlooking heterogeneous data that may require distinct feature representations. RTL discovers specialized, routed subnetworks and achieves superior or competitive performance across multiple settings with a compact parameter footprint.

  • Problem

    Nearly all Lottery Ticket Hypothesis methods assume one universal subnetwork across inputs, overlooking heterogeneous data that may require distinct feature representations.

  • Method

    RTL jointly discovers multiple specialized, sparse subnetworks from a shared initialization, using mask-based context routing without auxiliary routing networks or additional parameters.

  • Results

    Up to an order of magnitude fewer parameters than independent models, RTL outperforms single-mask and multi-model pruning baselines on CIFAR-10 and scales to CIFAR-100 specialization.

  • Takeaways & Limitations

    Mask similarity provides a label-free indicator of subnetwork collapse, while adaptive pruning aligns model structure with heterogeneous data and supports compact, modular models.

  • Takeaways & Limitations

    RTL favors recall over precision, so downstream calibration or selective filtering may be needed to improve precision.

Abstract

from arXiv · show

In pruning, the Lottery Ticket Hypothesis posits that large networks contain sparse subnetworks, or winning tickets, that can be trained in isolation to match the performance of their dense counterparts. However, most existing approaches assume a single universal winning ticket shared across all inputs, ignoring the inherent heterogeneity of real-world data. In this work, we propose Routing the Lottery (RTL), an adaptive pruning framework that discovers multiple specialized subnetworks, called adaptive tickets, each tailored to a class, semantic cluster, or environmental condition. Across diverse datasets and tasks, RTL consistently outperforms single- and multi-model baselines in balanced accuracy and recall, while using up to 10 times fewer parameters than independent models and exhibiting semantically aligned. Furthermore, we identify subnetwork collapse, a performance drop under aggressive pruning, and introduce a subnetwork similarity score that enables label-free diagnosis of oversparsification. Overall, our results recast pruning as a mechanism for aligning model structure with data heterogeneity, paving the way toward more modular and context-aware deep learning.

1. Introduction

The introduction argues that conventional pruning is hindered by deployment inefficiency and universal masks that overlook heterogeneous data. RTL addresses this gap by routing among specialized, shared-backbone subnetworks and improves performance across image and speech tasks.

  • Motivation: Modern deep networks often require billions of parameters and hundreds of gigaflops per inference, limiting deployment in resource-constrained or real-time settings.Reducing model complexity is presented as important for practical deployment and for understanding generalization.
  • Problem: Most Lottery Ticket Hypothesis methods use one universal sparse mask, potentially sacrificing performance by forcing heterogeneous data through a rigid architecture.Distinct classes, clusters, or environmental conditions may rely on different feature representations.
  • RTL framework: Routing the Lottery discovers multiple adaptive tickets tailored to distinct data subsets, aligning sparsity with data structure rather than enforcing uniform compression.The framework allocates representational capacity heterogeneously across the input space.
  • RTL framework: RTL jointly learns data-adapted sparse subnetworks from one dense initialization, preserving parameter sharing through a compact backbone and mask-based routing.This enables context-aware inference while maintaining shared parameters.
  • Empirical validation: Class-specific RTL subnetworks outperform single-mask and multi-model pruning baselines on CIFAR-10 while using up to an order of magnitude fewer parameters than independent models.On CIFAR-100, RTL also scales to specialize across more classes.
  • Empirical validation: On speech enhancement, subnetworks specialized for acoustic environments achieve higher SI-SNRi than universal or independent baselines.This validates RTL on a real-world task beyond image classification.

2. Related Work

Related work progresses from classical weight-pruning methods and lottery-ticket discovery toward dynamic sparsity and input-dependent subnetworks. Conditional computation improves scalability but typically incurs complex routing, auxiliary parameters, and substantial computation.

  • Classical pruning: Classical pruning methods include sensitivity-based approaches and unstructured removal of non-relevant weights without eliminating entire layers.Examples include Optimal Brain Damage, Optimal Brain Surgeon, meaningful-connection learning, and energy-based pruning.
  • Lottery tickets: Iterative Magnitude Pruning identifies winning tickets but remains computationally intensive because it requires repeated training cycles.Subsequent work studies mask information, SGD, initialization, learning-rate rewinding, and theoretical explanations.
  • Dynamic sparsity: Dynamic sparse training evolves sparsity patterns during training by pruning and regrowing connections online, relaxing the fixed-mask assumption.SET, SNFS, and RigL implicitly acknowledge that different inputs may activate different network pathways.
  • Conditional computation: Conditional computation methods such as MoE and conditional convolutions activate input-dependent subnetworks to improve scalability.These approaches typically require complex routing mechanisms, large auxiliary parameter sets, and substantial compute budgets.

3. Method

RTL extends lottery-ticket pruning by learning multiple specialized subnetworks for distinct data subsets rather than one universal mask. It partitions data, extracts subset-specific masks under sparsity constraints, and jointly retrains the resulting tickets while preserving mask structure and limiting interference.

  • Method overview: RTL discovers distinct subnetworks, each tailored to a specific data cluster, instead of seeking a single universal winning ticket.The framework extends the Lottery Ticket Hypothesis to heterogeneous data by learning specialized masks for different subsets.
  • Adaptive pruning objective: The dataset is partitioned into K subsets corresponding to classes or clusters, using supervised labels or unsupervised clustering.RTL remains agnostic to how subsets are defined, supporting manual labeling and automatic clustering across tasks and modalities.
  • Adaptive pruning objective: For each subset D_k, RTL learns a dedicated mask m_k and jointly optimizes the mask set under a sparsity constraint.Each mask defines a specialized subnetwork, with sparsity controlling the maximum number of non-zero parameters per subnetwork.
  • Adaptive mask extraction: Each pruning iteration trains on a subset, removes the lowest-magnitude weights by fraction p, resets remaining weights to θ_0, and repeats until density reaches s.The resulting masks form M = {m_1, . . . , m_K}, and each specialized subnetwork is f(x; m_k ⊙θ_0).
  • Joint retraining: Joint retraining preserves mask topology and updates only active weights, while masked gradients prevent interference between subnetworks and help avoid catastrophic forgetting and collapse.Balanced batches ensure each subnetwork receives the same number of gradient updates per epoch despite differing subset sizes.

4. Experiments

The experiments progressively evaluate RTL from controlled class-specific pruning to cluster-, region-, and environment-specific applications, while analyzing subnetwork overlap and semantic alignment. Across tasks, comparisons control architecture, sparsity, and training budget to isolate specialization effects.

  • Experimental overview: Experiments span class-specific subnetworks on CIFAR-10, semantic clusters on CIFAR-100, within-image INR regions, heterogeneous speech environments, and subnetwork alignment analysis.The sequence moves from controlled settings toward real-world applications and structural analysis.
  • CIFAR-10: On CIFAR-10, RTL assigns each of 10 classes a subnetwork and is compared with single-model and independently pruned multiple-model IMP baselines.All methods share the backbone, sparsity targets, and pruning budget; evaluation covers balanced accuracy, precision, recall, and parameter count at 25%, 50%, and 75% sparsity.
  • CIFAR-100: On CIFAR-100, RTL assigns subnetworks to 8 unsupervised coarse semantic clusters, testing specialization under fine-grained, overlapping, and imperfectly aligned class boundaries.RTL is compared with the same two IMP baselines using matched metrics and sparsity levels.
  • Implicit neural representations: For INRs, RTL specializes subnetworks to semantic regions within each of 10 ADE20K images, reconstructing RGB values from continuous pixel coordinates.Segmentation masks define region classes, while the baseline conditions one network on region identity through class embeddings; reconstruction is measured by PSNR.
  • Speech enhancement: For speech enhancement, RTL learns one subnetwork for each of three acoustic scenes and is compared with shared and independently pruned IMP models using SI-SNRi.The scenes are indoor, outdoor, and transportation, and all methods operate under identical sparsity and computational constraints.
  • Subnetwork analysis: CIFAR-10 and CIFAR-100 analyses quantify mask overlap across sparsity levels and relate excessive overlap to balanced-accuracy degradation, termed subnetwork collapse.CIFAR-10 additionally compares structural similarity with WordNet-derived semantic distances between class labels.

5. Results

RTL consistently outperforms single- and multi-model pruning baselines across classification, image reconstruction, and speech enhancement, while using compact specialized subnetworks. Its results also show that excessive mask overlap predicts subnetwork collapse and that semantic structure emerges during pruning and training.

  • CIFAR-10: On CIFAR-10, RTL achieves the highest balanced accuracy and recall across sparsity levels, including balanced accuracy 0.781 at 25% sparsity versus 0.711 for single-model IMP and 0.712 for multi-model IMP.At 75% sparsity, balanced accuracy remains 0.772 versus 0.760 for multi-model IMP; recall is 0.821 / 0.810 / 0.816 across sparsity levels.
  • CIFAR-10: RTL’s CIFAR-10 precision is lower, at 0.257-0.282 versus 0.478-0.515 for single-model IMP, reflecting a recall-prioritizing sensitivity–discrimination trade-off.The passage identifies thresholding or ensemble methods as ways to recover precision downstream.
  • CIFAR-100: On CIFAR-100 semantic clusters, RTL leads balanced accuracy at 0.765 / 0.751 / 0.759 and recall at 0.764 / 0.729 / 0.754 across 25% / 50% / 75% sparsity.At 75% sparsity, RTL remains ahead despite noisy cluster boundaries; its parameters are 108K / 76K / 40K versus 944K / 629K / 314K for multi-model IMP.
  • Image reconstruction: For coordinate-based image reconstruction, RTL exceeds single-model IMP by almost 3 dB at 25% sparsity and by 2.53 dB and 2.18 dB at 50% and 75%.RTL degrades more gracefully as sparsity increases, supporting region-specific subnetworks for semantically distinct image regions; per-image results remain consistent across diverse scenes.
  • Speech enhancement: Across indoor, outdoor, and transportation environments, RTL attains SI-SNRi of 7.248, 7.178, and 6.992 at 25%, 50%, and 75% sparsity, outperforming both IMP baselines.RTL uses 32.0K / 22.8K / 12.3K parameters, while multi-model IMP uses 84.1K versus 32.0K at 25% sparsity; independent pruning without shared initialization or joint retraining underperforms.
  • Subnetwork collapse and semantic structure: Mask similarity provides a label-free collapse signal: IoU spikes precede or coincide with sharp accuracy drops once pruning forces excessive subnetwork overlap.Collapse occurs at slightly higher sparsity thresholds on CIFAR-100 than CIFAR-10, while moderate pruning can improve early-layer organization and later training increases middle- and deep-layer correlations.

6. Conclusion · A. Vision Model and Training Setup

Routing the Lottery discovers specialized adaptive tickets that preserve compactness while aligning subnetworks with heterogeneous data. Its effectiveness depends on structural diversity, remains robust across partition quality, and may benefit from downstream precision calibration.

  • 6. Conclusion: RTL discovers multiple specialized subnetworks, called adaptive tickets, instead of a single universal winning ticket.Specialization emerges when subnetworks diverge in response to data heterogeneity.
  • 6. Conclusion: RTL achieves superior or competitive performance across multiple settings while maintaining a compact parameter footprint.The framework positions adaptive pruning as both an efficient and performance-preserving alternative to universal subnetworks.
  • 6. Conclusion: Mask similarity provides a reliable, label-free indicator of subnetwork collapse by measuring whether structural diversity is preserved.The analysis identifies preservation of distinct masks, rather than raw capacity alone, as critical to avoiding collapse.
  • 6. Conclusion: RTL remains effective with clean class boundaries, noisy semantic clustering, and real-world applications, demonstrating robustness to imperfect data partitions.The passage specifically contrasts CIFAR-10’s clean boundaries with CIFAR-100’s noisy semantic clustering before noting real-world generalization.
  • 6. Conclusion: RTL subnetworks tend to favor recall over precision, increasing sensitivity to weak class- or environment-specific signals.This behavior benefits routing-based inference but may require downstream calibration or selective filtering to improve precision.
  • 6. Conclusion: Adaptive pruning dynamically allocates representational capacity in alignment with the intrinsic structure of processed data.The resulting models are framed as more modular, efficient, and interpretable.

A.1. Model architecture and pruning scope · A.2. Optimization and training protocol · A.3. Pruning schedule and compute

The CIFAR experiments use a truncated GhostNet with masked convolutions as the pruning scope, while dense normalization, squeeze-and-excitation, and classifier components are retained. Training uses Adam with balanced-comparison protocols, fixed per-epoch pruning, and single-GPU execution with reported runtime differences across methods.

  • A.1. Model architecture and pruning scope: GhostNet serves as the CIFAR-10 and CIFAR-100 backbone, retaining the convolutional stem and first nine Ghost bottleneck blocks.Deeper stages of the original network are discarded to control model capacity and isolate adaptive pruning.
  • A.1. Model architecture and pruning scope: Masked convolutions define the pruning scope, while batch normalization, squeeze-and-excitation modules, and the final classifier remain dense.The configuration contains approximately 126K prunable parameters.
  • A.1. Model architecture and pruning scope: The retained backbone follows the standard Ghost bottleneck design with depthwise convolutions, pointwise expansions, optional downsampling, and progressive spatial reduction.Selected bottlenecks use strided depthwise convolutions, and Table 4 documents channel dimensions and downsampling stages for reproducibility.
  • A.2. Optimization and training protocol: All models use Adam with a learning rate of 1e−4 and no weight decay, with pruning and joint retraining phases lasting 10 epochs each.Batch sizes are 320 for CIFAR-10 and 256 for CIFAR-100.
  • A.2. Optimization and training protocol: RTL, single-model IMP, and multi-model IMP share the same random Kaiming initialization for fair comparison.RTL and multi-model IMP use balanced binary-classifier mini-batches, while single-model IMP trains for 20 epochs to match total data exposure.
  • A.3. Pruning schedule and compute: Pruning removes 4,096 weights per epoch for each subnetwork until the target sparsity is reached.All experiments run on a single NVIDIA H100 GPU.
  • A.3. Pruning schedule and compute: 6 hours is the approximate runtime for full RTL and multi-model IMP runs, compared with roughly 45 minutes for the single-model IMP baseline.These runtimes are reported for execution on a single NVIDIA H100 GPU.

B. Semantic Clustering Procedure … C.6. Evaluation protocol

The paper constructs semantic CIFAR-100 subsets with fixed, text-based clustering and evaluates RTL on implicit neural representations using independently trained ADE20K images. The protocol specifies Fourier-coordinate inputs, MLP architectures, iterative prune-and-rewind training, matched initialization, and PSNR evaluation.

  • B. Semantic Clustering Procedure: RTL forms 8 CIFAR-100 semantic clusters from class-name text embeddings reduced with UMAP and clustered using HDBSCAN.The clusters are semantically coherent but intentionally imperfectly aligned with visual similarity.
  • B. Semantic Clustering Procedure: Clustering hyperparameters remain fixed rather than tuned to downstream performance, while RTL remains effective under non-ideal specialization boundaries.This setting tests adaptive pruning beyond strictly class-aligned partitions.
  • C. Implicit Neural Representation: The INR task maps continuous pixel coordinates (x, y) to RGB values, reconstructing each image through point-wise network evaluation.Experiments use 10 diverse ADE20K training images, each modeled independently.
  • C.1. Task setup: Within each image, semantic segmentation masks merge all instances of the same object category into one class.Small unassigned or noisy boundary regions are reassigned to neighboring semantic classes to stabilize training.
  • C.2. Positional encoding: Input coordinates are encoded with Fourier features, and this encoded representation is the RTL model’s sole input.The positional encoding follows standard INR practice.
  • C.3. Model architectures: RTL and baseline models are five-layer ReLU MLPs, with RTL using dimensions [[128, 34], [128, 128], [128, 128], [128, 128], [3, 128]].The baseline adds a learned class embedding, changing only the first-layer dimensionality to [128, 50].
  • C.4. Training and pruning protocol: Training uses Adam at learning rate 0.01 for 10,000 steps, followed by pruning 4,096 weights per iteration and rewinding remaining weights after each step.Prune-and-rewind repeats until the target sparsity is reached.
  • C.5. Initialization and fairness considerations: Fair initialization removes the baseline’s class-embedding weights to form the RTL initialization, keeps shared parameters identical, and uses one seed across images.Reconstruction quality is reported with PSNR averaged over all pixels and images, with per-image values additionally provided in the appendix.

C.7. Per-Image INR Results

RTL consistently outperforms the single-model IMP baseline in per-image PSNR across 10 ADE20K images and sparsity levels. At high sparsity, reconstruction quality drops sharply alongside mask-similarity collapse, supporting region-specific pruning masks.

  • Per-image reconstruction performance: RTL achieves higher per-image PSNR than the single-model IMP baseline across all 10 ADE20K images and at 25%, 50%, and 75% sparsity.The results show that aggregate improvements are consistent across diverse scene content rather than driven by a small subset of images.
  • Per-image reconstruction performance: RTL’s advantage is particularly pronounced at higher sparsity, where specialized subnetworks preserve region-specific parameters and produce more stable degradation than a single global pruning mask.A global mask causes larger reconstruction errors across semantically heterogeneous regions.
  • Implications: Per-image trends mirror averaged main-text results and support maintaining distinct, region-specific pruning masks for high-sparsity INR reconstruction.The per-image analysis shows that adaptive-pruning benefits are not explained by outliers or dataset bias.
  • PSNR and mask similarity: Across images, PSNR decreases gradually with sparsity before dropping sharply, while pruning-mask similarity remains stable at low-to-moderate sparsity and rapidly declines beyond a critical threshold.These trends persist despite substantial differences in scene content and appearance.
  • PSNR and mask similarity: The sharp PSNR decline coincides with mask-similarity collapse, indicating that semantic-region subnetworks interfere substantially once their masks overlap or become unstable.At very high sparsity, mask similarity approaches zero and reconstruction quality becomes poor.

C.8. INR subnetwork similarity vs. reconstruction quality … D.4. Optimization and runtime

RTL preserves specialized subnetworks under aggressive pruning, improving INR reconstruction quality and qualitative speech or image outputs while enabling label-free diagnosis of subnetwork collapse. The speech-enhancement setup uses environment-specific subnetworks within a compact U-Net and reports approximately 8–10-hour training runtimes.

  • C.8. INR subnetwork similarity vs. reconstruction quality: Increasing mask similarity coincides with sharp PSNR degradation, indicating subnetwork collapse under aggressive pruning across four- and fifteen-region ADE20K images.The trend holds across all fifteen semantic regions, and mean-centered fluctuations in PSNR closely track mask-similarity changes.
  • C.8. INR subnetwork similarity vs. reconstruction quality: ρ ≈0.982: At 50% sparsity, PSNR and mask similarity show a strong monotonic relationship across regions within a single ADE20K image.This supports subnetwork similarity as a strong predictor of INR reconstruction quality.
  • C.8. INR subnetwork similarity vs. reconstruction quality: Mask similarity provides a reliable, label-free diagnostic for oversparsification and impending performance collapse in coordinate-based representations.The analysis links preserved subnetwork distinctiveness to fine-grained, within-image semantic specialization.
  • C.9. Qualitative INR Reconstructions under High Sparsity: At 50% sparsity, RTL preserves finer details and sharper object boundaries than single-mask IMP, especially in textured regions and along semantic edges.Both methods recover overall scene structure, but IMP shows mild blurring and contrast loss.
  • D.1. Dataset construction: Environmental noise is grouped into indoor, outdoor, and transportation scenes, with each scene defining a distinct specialized subnetwork for speech enhancement.Clean speech comes from DNS Challenge 2020 and environmental noise from TAU Urban Acoustic Scenes 2020.
  • D.2. Model architecture: The speech-enhancement model is a compact five-stage encoder–decoder U-Net using complex STFT inputs and predicting a two-channel complex ratio mask.Masked convolutions, time-only downsampling, ELU activation, batch normalization, and skip connections define the architecture.
  • D.3. Signal processing and loss: 10-second waveforms produce [2, 626, 513] spectrograms, and training uses weighted source-to-distortion ratio loss; pruning and retraining each run for 10 epochs at learning rate 1e−4.Experiments use Adam without weight decay, identical initial weights, and a single NVIDIA H100 GPU; multi-model IMP takes approximately 10 hours versus about 8 hours for single-model IMP.

D.5. Qualitative Analysis of Speech Enhancement

At 50% weight sparsity, RTL qualitatively produces cleaner and more faithful speech enhancement across three acoustic environments than the single-mask IMP baseline. Its environment-specialized subnetworks better preserve speech structure under challenging, heterogeneous noise conditions, consistent with reported SI-SNRi improvements.

  • Qualitative comparison: At 50% weight sparsity, Fig. 12 compares RTL and IMP outputs across three acoustic environments using noisy inputs, clean targets, and enhanced spectrograms.Each row represents an environment class, while columns show the noisy input, clean target, and model outputs.
  • Qualitative comparison: RTL more faithfully reconstructs clean-speech time-frequency structure, with clearer harmonics, better transient preservation, and stronger suppression of noise-dominated regions.IMP instead shows residual noise, smeared harmonics, and reduced speech-background contrast, especially at mid and high frequencies.
  • Challenging conditions: In challenging conditions dominated by environment-specific noise, RTL recovers speech with higher temporal coherence and sharper spectral detail than IMP.Environment-specialized subnetworks adapt to acoustic scenes, whereas the single-mask IMP model struggles to balance denoising across heterogeneous conditions.
  • Interpretation: The visual comparisons align with quantitative SI-SNRi improvements and support environment-aligned adaptive pruning over a single global sparse model.These results are presented as evidence that adaptive pruning yields more effective representations.

E. Subnetwork Similarity Analysis … E.3. Semantic alignment

The analysis quantifies pairwise pruning-mask similarity across sparsity levels, layers, and subnetworks, then relates structural overlap to performance and semantic relatedness. It uses balanced accuracy to examine subnetwork collapse and CIFAR-10 WordNet distances to test semantic alignment.

  • E.1. Mask similarity: Pairwise similarity between binary pruning masks is computed at multiple sparsity levels using the Jaccard coefficient.The comparison is defined for masks Mi and Mj.
  • E.1. Mask similarity: Jaccard intersection and union are computed element-wise across all prunable parameters.
  • E.1. Mask similarity: Similarity is evaluated globally and per layer to measure how mask overlap changes across network depth.
  • E.2. Collapse analysis: Balanced accuracy at multiple sparsity levels is correlated with each subnetwork’s average pairwise mask similarity to assess collapse-related performance degradation.Excessive overlap is treated as subnetwork collapse.
  • E.3. Semantic alignment: CIFAR-10 semantic distances between class labels are computed with WordNet path similarity and compared with corresponding mask similarities.The comparison tests whether conceptually related classes share pruning structure.
  • E.3. Semantic alignment: Aligned semantic and structural similarity matrices reveal whether related CIFAR-10 classes have similar pruning structures.
Loading 2601.22141v1…