Switchable Output Layer (SOL) in Neural Networks

• Switchable Output Layer (SOL) is a design pattern that assigns independent classification heads to each sub-net in a shared network, reducing gradient conflicts and representational interference.
• It decouples logit learning within once-for-all training frameworks, enabling each sub-net to optimize decision boundaries for improved accuracy, adversarial robustness, and calibration.
• SOL is adaptable across domains, demonstrated in neural network classification and ultrasound beamforming via AdaIN, which streamlines deployment without altering runtime cost.

Switchable Output Layer (SOL) denotes a switchable component that changes the output behavior of a largely shared neural network. In once-for-all training, SOL is defined as a separate, unshared classification head for each sub-net in a shared super-network, with the backbone kept shared; in the SOLAR formulation, this decouples logit learning across sub-nets and is reported to improve accuracy, adversarial robustness, and calibration without changing the backbone or incurring runtime cost for any single sub-net (Tareen et al., 20 Sep 2025). In a distinct ultrasound beamforming setting, the switch is implemented through Adaptive Instance Normalization (AdaIN) at the bottleneck, allowing one network to produce DAS, speckle-reduced, deconvolved, and deconvolved+despeckled outputs by changing a compact code (Khan et al., 2020). Across these uses, the common idea is selective switching of a small output-defining mechanism while preserving a shared feature-extraction path.

1. Core definition and motivation

In the SOLAR formulation, SOL arises from a specific bottleneck in Once-for-All (OFA) training. OFA training enables a single super-net to generate multiple sub-nets tailored to diverse deployment scenarios, supporting flexible trade-offs among accuracy, robustness, and model-size without retraining. Standard practice shares not only the backbone but also a single output layer across sub-nets of different capacities. The paper identifies this last-layer sharing as a source of excessive parameter sharing, gradient conflict, representational interference, degraded calibration, and reduced overall performance, especially as the number of supported sub-nets increases (Tareen et al., 20 Sep 2025).

SOL addresses this by assigning each sub-net its own classification head while leaving the backbone and training procedures unchanged. Even when switchable BN or SDBN layers separate feature statistics, all sub-nets still collide at the final linear classifier; SOL specifically removes that collision at the logit layer. The resulting mechanism is narrow in scope but targeted: the backbone remains shared, and only the terminal classifier becomes switchable.

The beamforming literature uses the same switchability principle differently. In “Switchable Deep Beamformer,” a single neural beamforming network produces multiple output types by changing a compact task code, with the switch implemented through AdaIN layers. The details explicitly identify one AdaIN layer at the bottleneck as the “Switchable Output Layer (SOL)” in that implementation. The motivation is analogous in systems terms: replacing multiple task-specific models with a single shared model plus a small switchable component reduces storage and simplifies deployment (Khan et al., 2020).

2. Architecture and formalism in once-for-all training

Let the shared backbone $f_\theta$ produce a feature vector $h \in \mathbb{R}^d$ from an input $x$. For $K$ supported sub-nets, SOL maintains $K$ independent classification heads parameterized by $$\{W_k \in \mathbb{R}^{C \times d}, b_k \in \mathbb{R}^C\}$$, one per sub-net $k$. When sub-net $k$ is active, its logits are

$z^{(k)} = W_k h + b_k.$

With softmax

$$p^{(k)}_c = \frac{\exp(z^{(k)}_c)}{\sum_{j=1}^C \exp(z^{(k)}_j)},$$

the cross-entropy for a one-hot target $h \in \mathbb{R}^d$0 is

$h \in \mathbb{R}^d$1

For adversarial training, the reported hybrid loss uses a convex combination of clean and adversarial losses weighted by $h \in \mathbb{R}^d$2:

$h \in \mathbb{R}^d$3

where

$h \in \mathbb{R}^d$4

and PGD-7 with $h \in \mathbb{R}^d$5 and step size $h \in \mathbb{R}^d$6 is used for inner maximization. If multiple sub-nets are evaluated per batch, the total loss is

$h \in \mathbb{R}^d$7

where $h \in \mathbb{R}^d$8 is the set of sub-nets sampled or activated in the step. In OATS-SOL experiments with 8 widths, $h \in \mathbb{R}^d$9 typically includes all width fractions $x$0 (Tareen et al., 20 Sep 2025).

This formulation makes the role of SOL precise. The shared encoder still determines the feature vector $x$1, but the final class-score map is sub-net-specific. The paper argues that this lets each width optimize decision boundaries consistent with its capacity and feature distribution, reducing interference and improving optimization. SOL is also described as agnostic to the learning objective and compatible with standard CE, Kullback–Leibler distillation, TRADES, MMA, and adversarial distillation, although the main robustness experiments use CE+PGD-based hybrid losses.

3. Integration with OATS and SNN frameworks

SOL is reported as integrating seamlessly with Slimmable Neural Networks (SNNs) and Once-for-All Adversarial Training and Slimming (OATS). Convolutional and linear backbone layers remain shared; SBN or SDBN remain per-sub-net; only the output heads become switchable. The training loop is described conceptually as follows: sample a batch $x$2 and sample $x$3 from the set $x$4 in OATS; for each supported width $x$5, activate sub-net $x$6 by gating channels and selecting its SBN or SDBN parameters; forward $x$7 through the shared backbone to obtain $x$8 and then through the active head $x$9; generate an adversarial example $K$0 with PGD-7; compute the hybrid loss; backpropagate and accumulate gradients; and finally update the shared backbone and active head parameters (Tareen et al., 20 Sep 2025).

The reported experiments use four super-net backbones: ResNet-34, WideResNet-16-8, WideResNet-40-2, and MobileNetV2. No architectural constraints are imposed beyond having a shared encoder and a terminal linear layer. Width-adaptive training uses SBN in SNN and SDBN in OATS to maintain separate normalization statistics per width and, for OATS, per data type. FiLM layers used in OATS remain compatible.

The hyperparameter settings reported in the paper are specific. The datasets are SVHN, CIFAR-10, STL-10, CIFAR-100, and TinyImageNet, with 10, 10, 10, 100, and 200 classes respectively. OATS-SOL typically uses 8 widths for SVHN and CIFAR-10, and 4 widths for STL-10 and CIFAR-100. SNN-SOL experiments include 8 and 16 widths, and TinyImageNet runs use 8 widths. Optimization uses SGD with momentum 0.9, a cosine annealing scheduler, learning rates from $K$1, batch size 64 for STL-10 and TinyImageNet and 128 for the other datasets, and training for 40 epochs on SVHN in OATS-SOL, 120 epochs on CIFAR-10, CIFAR-100, and STL-10, and 140 epochs on TinyImageNet in SNN-SOL. The adversarial setting remains PGD-7 with $K$2 $K$3 and step size $K$4, with uniform sampling of $K$5 from the specified set during training and validation; $K$6 is tunable at inference to select an accuracy–robustness point.

4. Calibration, robustness, and empirical behavior

The SOLAR paper motivates SOL partly through calibration. It states that the last layer maps backbone features into class scores and implicitly defines margins; when heterogeneous sub-nets produce different feature distributions and capacities, a shared head forces a single set of margins, leading to miscalibration and suboptimal robustness. Per-sub-net heads allow each width to learn an appropriate scaling and orientation of decision boundaries and logits. The discussion of calibration in the main text is qualitative rather than metric-centered, although the paper provides the standard definition of Expected Calibration Error,

$K$7

and notes negative log-likelihood and Brier score as other common metrics. Robustness is evaluated under adversarial perturbations generated by PGD-7 at test time with $K$8 norm constraints $K$9, and reported “robustness” is PGD accuracy (Tareen et al., 20 Sep 2025).

The paper reports improvements over OATS on five datasets. On SVHN with WideResNet-16-8 and 8 sub-nets, SOLAR improves accuracy of sub-nets by up to 1.26% and robustness by up to 9.01%; the smallest sub-net reaches 94.01% accuracy and 53.08% robustness, surpassing OATS by 0.57% accuracy and 1.57% robustness, with model size $K$0 KB. On CIFAR-10 with ResNet-34 and 8 sub-nets, the reported gains are up to 4.71% in accuracy and 7.71% in robustness. On STL-10 with WideResNet-40-2 and 4 sub-nets, the reported gains are up to 1.67% in accuracy and 2.72% in robustness. On CIFAR-100 with ResNet-34 and 4 sub-nets, the reported gains are up to 1.76% in accuracy and 1.26% in robustness (Tareen et al., 20 Sep 2025).

Against SNN baselines, the reported gains are in standard accuracy. On CIFAR-10 with 16 sub-nets, SNN-SOL improves per-sub-net accuracy by up to 2.67% for MobileNetV2 at 2/16 width, 2.43% for WideResNet-16-8 at 3/16 width, and 0.98% for ResNet-34 at 11/16 width. On CIFAR-100 with 8 sub-nets, the reported gains reach 1.65% for ResNet-34 at 8/8, 0.46% for WideResNet-16-8 at 7/8, and 0.96% for MobileNetV2 at 8/8; the paper also notes that some very small widths show negligible or slightly negative changes. On TinyImageNet with 8 sub-nets, the reported gains are up to 2.93% for ResNet-34 at 8/8, 2.34% for WideResNet-16-8 at 7/8, and 1.35% for MobileNetV2 at 3/8. Post-search fine-tuning preserves the advantage of SOL sub-nets, with example gains after fine-tuning of +2.42% on CIFAR-10 for MobileNetV2 at 2/16 and +2.19% on TinyImageNet for ResNet-34 at 5/8. This suggests that unshared heads can yield stronger initial representations and more generalizable minima.

5. Computational cost, deployment semantics, and common misunderstandings

The parameter overhead of SOL is explicit. Each head adds $K$1 parameters, and the total super-net overhead for $K$2 heads is $K$3. The reported point, however, is deployment-specific: the parameters used by any single deployed sub-net are unchanged relative to the baseline because the deployed sub-net carries only its own head, and other heads can be discarded. In the super-net, empirical parameter increases with 8 sub-nets on CIFAR-10 are 0.08% for ResNet-34, 0.16% for WideResNet-16-8, 0.97% for WideResNet-40-2, and 1.91% for MobileNetV2. Even with 64 sub-nets, the reported increases are 0.74%, 1.45%, 6.88%, and 12.32% for the same backbones (Tareen et al., 20 Sep 2025).

For FLOPs and runtime, the paper states that for a given sub-net the cost is identical to the baseline because the output layer’s shape and operations per sub-net do not change when “switching” replaces “slimming”:

$K$4

Measured FLOPs for matched sub-nets are reported as confirming no overhead in training or inference. The practical interpretation is that SOL increases super-net storage but does not change the execution cost of a selected sub-net.

A recurring misunderstanding is to treat SOL as a backbone modification. In SOLAR, the shared backbone remains unchanged and fixed during a given forward pass; only the head is switched when the active sub-net changes. Another misunderstanding is to equate super-net storage overhead with deployment overhead. The paper explicitly distinguishes these: the extra heads live in the super-net checkpoint, while a deployed sub-net retains only its own head. The implementation is described as purely architectural, compatible with mixed-precision, standard serialization, and model export, and straightforward in checkpoint management because the heads are small relative to the backbones.

6. AdaIN-based switchable output in deep beamforming

In “Switchable Deep Beamformer,” the switchable component is implemented through Adaptive Instance Normalization at the bottleneck of a single beamforming network rather than through multiple terminal classifiers. The model is designed to produce four outputs—DAS, speckle removal, deconvolution, and deconvolution+despeckle—using one shared network with a simple switch (Khan et al., 2020).

The baseline backbone consists of nine convolutional blocks: blue blocks contain two Conv layers with ReLU, yellow blocks contain Conv plus LeakyReLU, all Conv layers use stride $K$5 and kernel $K$6, and the final layer uses kernel $K$7 and stride $K$8 to project features to the $K$9 output. The input is a multi-depth, multi-line slice of channel data around depth $$\{W_k \in \mathbb{R}^{C \times d}, b_k \in \mathbb{R}^C\}$$0,

$$\{W_k \in \mathbb{R}^{C \times d}, b_k \in \mathbb{R}^C\}$$1

The bottleneck feature map $$\{W_k \in \mathbb{R}^{C \times d}, b_k \in \mathbb{R}^C\}$$2 is transformed by AdaIN conditioned on a task code. For channel $$\{W_k \in \mathbb{R}^{C \times d}, b_k \in \mathbb{R}^C\}$$3,

$$\{W_k \in \mathbb{R}^{C \times d}, b_k \in \mathbb{R}^C\}$$4

and the task-conditioned transform is

$$\{W_k \in \mathbb{R}^{C \times d}, b_k \in \mathbb{R}^C\}$$5

A small MLP maps a scalar task code to the channelwise mean and variance vectors $$\{W_k \in \mathbb{R}^{C \times d}, b_k \in \mathbb{R}^C\}$$6 and $$\{W_k \in \mathbb{R}^{C \times d}, b_k \in \mathbb{R}^C\}$$7; the mean head uses linear activation and the variance head uses ReLU. The code generator is trained jointly with the beamformer, but at inference it can be discarded and replaced by stored AdaIN code vectors.

The task targets are also defined explicitly. DAS mode learns the log-compressed absolute Hilbert envelope of the DAS RF sum:

$$\{W_k \in \mathbb{R}^{C \times d}, b_k \in \mathbb{R}^C\}$$8

Deconvolution mode is based on the forward model $$\{W_k \in \mathbb{R}^{C \times d}, b_k \in \mathbb{R}^C\}$$9 and the sparse blind deconvolution problem

$k$0

with target $k$1. Speckle-removal targets apply NLLR filtering to the log-absolute Hilbert envelope, and the combined mode applies NLLR to the deconvolved, log-compressed image.

Training minimizes

$k$2

using L2 loss, Adam, learning rate $k$3 with exponential decay, 100 epochs, and early stopping at 69 epochs with patience 20. The dataset comprises carotid, thyroid, forearm, calf, and phantom scans acquired with E-CUBE 12R and an L3-12H linear array. The reported runtime is approximately 0.165 ms per depth plane, compared with approximately 4.8 ms for the DeepBF baseline. Quantitatively, the switchable model is reported to match or improve DAS GCNR, to exceed CNR and GCNR in Despeckle mode, and to yield sharper images in Deconv modes, with Deconv+Despeckle balancing sharpness and noise.

7. Limitations, generalization, and conceptual interpretation

The SOLAR paper lists concrete limitations. SOL increases super-net parameter storage proportional to $k$4, and gains are largest for smaller widths and larger $k$5. For some larger widths or certain datasets, including a few CIFAR-100 small-width cases, improvements can be modest or occasionally negative. If $k$6 is extremely large, head storage can become nontrivial, especially for backbones with large $k$7 and many classes, such as 200-class TinyImageNet (Tareen et al., 20 Sep 2025).

The beamforming implementation has a different set of limitations. Deconvolution mode can increase noise, reflected in lower CNR and GCNR despite improved CR and sharpness. Blind deconvolution depends on PSF estimation, so target quality can propagate errors into learning. The paper also notes that out-of-distribution shifts, such as different probes or plane-wave acquisitions, may require re-training or additional codes, and that incorrect or ambiguous codes produce mismatched statistics and suboptimal outputs (Khan et al., 2020).

The two lines of work support a common conceptual interpretation while remaining technically distinct. In SOLAR, the switchable component is an explicit family of terminal classifiers indexed by sub-net identity. In the beamforming system, the switchable component is a bottleneck AdaIN transformation whose code steers the network toward a desired output domain. This suggests that “Switchable Output Layer” is best understood as a design pattern rather than a single architecture: a mostly shared network is paired with a lightweight switchable module that resolves interference or multiplexes output behavior. The cited extensions are correspondingly broad. For SOLAR, the paper proposes scaling to ImageNet-1K, exploring normalization choices to further mitigate representational interference, and applying SOL in self-supervised and contrastive OFA settings. For the beamforming formulation, the paper discusses adding more tasks, exploring multi-layer conditioning, and replacing AdaIN with alternatives such as FiLM or conditional normalization.