# A Supernet-Only Framework for Federated Learning in Computationally Heterogeneous Scenarios
## 3 Method
### Reinterpret Approach
Secondly, the omission of certain layers alters the distribution of the aggregated representations, which consequently leads to a shift in the distribution of the final softmax output as follows: \(\text{Softmax}(\mathcal{N}(x))\). Although most networks are not entirely composed of residual units, the primary feature extraction within each stage is carried out by a sequence of residual units. Therefore, the transformation described by Equation (7) can be applied to each stage individually, with truncation performed within each stage accordingly. Assume the input to the \(s\)-th stage, which contains \(L_{s}\) residual units, is denoted as \(x_{s}\) and the output as \(y_{s}\). The result of a downsampling operation serves as the input \(x_{s,1}\) to the first residual unit. The input to the \(l\)-th residual unit \(\mathcal{N}_{s,l}\) within this stage is represented by \(x_{s,l}=\mathcal{N}_{s,l}^{\text{Pre}}(x_{s,1})\), where \(\mathcal{N}_{s,l}^{\text{Pre}}\) denotes the preceding net of the \(l\)-th layer within stage \(s\). Therefore, the featuring of this stage can be represented as follows:  
\[y_{s}=\sum_{l=1}^{L_{s}}\mathcal{W}_{s,l}\cdot\mathcal{N}_{s,l}^{\text{Raw}}( x_{s,l};\mathcal{P}_{s,l})+\mathcal{W}_{s,0}\cdot x_{s,1}, \tag{9}\]  
where \(\mathcal{W}_{s,l}\) and \(\mathcal{P}_{s,l}\) denote the weight and parameters of the \(l\)-th residual unit \(\mathcal{N}_{s,l}\) in stage \(s\), respectively. In manually designed networks such as ConvNeXt and Transformers, as well as in searched architectures like EfficientNet, the most common macro network structure typically consists of a preprocessing component, such as an Embedding layer or a Stem layer, followed by multiple stages. Each stage is composed of a downsampler followed by a feature extraction block, and the network ends with one or more detection heads. If each stage is rewritten in the form of Equation (9), while keeping all other components unchanged, the sequential network can be reinterpreted as a supernet for a NAS task. For the second issue, we chose to truncate the parameters of weights along with the network itself in order to reduce the magnitude of distribution shifts after dropping certain layers. Although the subsequent weights are discarded along with the truncated part of the network, the truncated weights still satisfy normalization and non-negativity, since only the weight parameters involved in subsequent computations are mapped to the probability simplex by the softmax function. Suppose in a certain training or inference scenario, stage \(s\) is truncated, retaining \(l_{s}\) residual units. Even if the weight parameters are not updated, the \(l_{s}+1\) weights that are used are recomputed as follows:  
\[\mathcal{W}_{s,l}=\frac{e^{\mathcal{W}_{s,l}}}{\sum_{j=0}^{l_{s}}e^{\mathcal{ W}_{s,l}}},\quad\text{where }0\leq l\leq l_{s}. \tag{10}\]  
The complete calculation of the truncated \(s\)-th stage is as follows:  
\[y_{s}=\sum_{l=1}^{l_{s}}\mathcal{W}_{s,l}\cdot\mathcal{N}_{s,l}^{\text{Raw}}( x_{s,l};\mathcal{P}_{s,l})+\mathcal{W}_{s,0}\cdot x_{s,1}\]  
\[\text{where}\quad \begin{cases}x_{s,1}=\text{Downsampler}_{s}(x_{s})\\ \mathcal{W}_{s,l}=\frac{e^{\mathcal{W}_{s,l}}}{\sum_{j=0}^{l_{s}}e^{\mathcal{ W}_{s,l}}},\quad 0\leq l\leq l_{s}\end{cases}  \tag{11}\]  
This process is illustrated in Figure 2, where each residual unit can be selectively retained up to \(l_{s}\) units per stage. By reinterpreting each stage in this way, while keeping other modules (e.g., Stem, Head) unchanged, any modern multi-stage residual network can be reformulated as a supernet suitable for NAS. As shown in Algorithm 1, the forward process of a truncated supernet can be represented as combination of several of the truncated stages described above.