# A Supernet-Only Framework for Federated Learning in Computationally Heterogeneous Scenarios
## 3 Method
### Reinterpret Approach
```
0: Input: \(x\), Truncation lengthes: \(L\), Network: \(\mathcal{N}=\{\text{Stem},\{\mathcal{N}_{s}\},\text{Head}\}\)

Each stage \(\mathcal{N}_{s}\) contains:

- Downsampler\({}_{s}\)

- Weight parameters \(\mathcal{W}_{s}^{\text{Raw}}=\{\mathcal{W}_{s,l}^{\text{Raw}}\}\quad l=1, \cdots,L_{s}\)

- Nonlinear transformations \(\mathcal{N}_{s}^{\text{Raw}}=\{\mathcal{N}_{s,l}^{\text{Raw}}\}\quad l=1, \cdots,L_{s}\)
0:\(y\)
1:\(x_{1}\leftarrow\text{Stem}(x)\);
2:for\(s=1\)to\(\text{len}(L)\)do
3:\(l_{s}\gets L[s]\);
4:\(\mathcal{W}_{s}[0:l_{s}]\leftarrow\text{Softmax}(\mathcal{W}_{s}^{\text{Raw}} [0:l_{s}])\);
5:\(y_{s,0}\leftarrow\text{Downsampler}_{s}(x_{s})\);
6:\(x_{s,1}\gets y_{s,0}\);
7:for\(l=1\)to\(l_{s}\)do
8:\(y_{s,l}\leftarrow\mathcal{N}_{s,l}^{\text{Raw}}(x_{s,l})\);
9:\(x_{s,l+1}\gets y_{s,l}+x_{s,l}\);
10:endfor
11:\(x_{s+1}\leftarrow\Sigma_{l=0}^{l_{s}}\mathcal{W}_{s,l}\cdot y_{s,l}\);
12:endfor
13:\(y\leftarrow\text{Head}(x_{s+1})\);
14:return\(y\);
```  
**Algorithm 1** Supernet forward algorithm  
The sub-network construction strategy can follow standard NAS- or pruning-based methods using architecture weights \(\mathcal{W}\). A simple instantiation used in our experiments is described in Section 4.1.3. As a supplementary illustration to Algorithm 1, we present a concrete forward pass of a two-stage supernet under client-specific truncation. Suppose Stage 1 contains two residual units, and Stage 2 contains ten. For a given client, the truncation policy selects \(L=[2,1]\), meaning that both residual units are active in Stage 1, and only one is active in Stage 2. The input \(x\) is first processed by the Stem module to obtain \(x_{1}\). Figure 2: Computation flow within stage \(\mathcal{N}_{s}\) of the supernet: The inputs are feature map \(x_{s}\) (from the previous stage), truncation length \(l_{s}\) (determined by client-side constraints), and raw weights \(\mathcal{W}_{s}\) (i.e., \(\mathcal{W}_{s}^{\text{Raw}}\)). \(D_{s}\), \(S\), and \(O_{s,l}\) denote the downsampler \(\text{Downsampler}_{s}\), the weight normalization function Softmax, and the nonlinear transformation \(\mathcal{N}_{s,l}^{\text{Raw}}\), respectively. The dashed box highlights a single residual unit, repeating \(l_{s}\) times. This stage corresponds to lines 4–11 in Algorithm 1. Stage 1. The input \(x_{1}\) is passed through the downsampler as follows:  
\[y_{1,0}=\text{Downsampler}_{1}(x_{1}),\quad x_{1,1}=y_{1,0}. \tag{12}\]  
The raw weights \(\mathcal{W}_{1}^{\text{Raw}}=[\mathcal{W}_{1,0}^{\text{Raw}},\mathcal{W}_{1,1} ^{\text{Raw}},\mathcal{W}_{1,2}^{\text{Raw}}]\) are normalized via softmax as follows:  
\[\mathcal{W}_{1,l}=\frac{e^{\mathcal{W}_{1,l}^{\text{Raw}}}}{\sum_{j=0}^{2}e^{ \mathcal{W}_{1,j}^{\text{Raw}}}},\quad l=0,1,2. \tag{13}\]  
Then the residual path is computed as follows:  
\[y_{1,1}=\mathcal{N}_{1,1}^{\text{Raw}}(x_{1,1}),\quad x_{1,2}=x_{1,1}+y_{1,1}, \tag{14}\]  
\[y_{1,2}=\mathcal{N}_{1,2}^{\text{Raw}}(x_{1,2}),\quad x_{1,3}=x_{1,2}+y_{1,2}. \tag{15}\]  
The output of Stage 1 is the weighted sum, as follows:  
\[x_{2}=\mathcal{W}_{1,0}\cdot y_{1,0}+\mathcal{W}_{1,1}\cdot y_{1,1}+\mathcal{ W}_{1,2}\cdot y_{1,2}. \tag{16}\]  
Stage 2. The input \(x_{2}\) is processed similarly.