UDC: Unified DNAS for Compressible TinyML Models for Neural Processing Units
Abstract
Deploying TinyML models on low-cost IoT hardware is very challenging, due to limited device memory capacity.
Neural processing unit (NPU) hardware address the memory challenge by using model compression to exploit weight quantization and sparsity to fit more parameters in the same footprint.
However, designing compressible neural networks (NNs) is challenging, as it expands the design space across which we must make balanced trade-offs.
This paper demonstrates Unified DNAS for Compressible ( UDC ) NNs, which explores a large search space to generate state-of-the-art compressible NNs for NPU.
ImageNet results show UDC networks are up to 3 .
35 × smaller (iso-accuracy) or 6 .
25% more accurate (iso-model size) than previous work.
Introduction
IoT applications demand TinyML models that fit on highly-constrained hardware (HW), with limited memory and compute power [14, 25, 30, 42].
Canonical TinyML tasks include visual wakewords, audio keywords, anomaly detection, speech enhancement and image classification [14, 13, 26, 9, 48].
Traditional microcontroller units (MCUs) are not well suited to meet the memory and compute challenges of TinyML, so HW vendors offer specialized processors for neural network (NN) inference, called neural processing units (NPUs) [2, 1].
MCU inference runtimes [7, 8] do not implement model compression, which is slow in software.
Hence MCUs do not benefit from sub 8-bit quantization or unstructured pruning [31].
In contrast, NPUs with HW model compression [2, 1] do benefit from both optimizations [3].
For example, the Arm Ethos-U55 NPU compiler (Vela [3]) encodes weight tensors using two components: 1) a binary mask marking non-zero elements (run-length and Golomb-Rice compressed), and 2) the non-zero values (Golomb-Rice compressed).
A model aggressively pruned to have many zeros and quantized to low bitwidths can be deployed in very little device memory.
At inference time, weights are then decompressed by the NPU on demand.
NPU compression is an enormous advantage for TinyML, where meager Flash memory limits model size and therefore accuracy [25, 48, 30, 42].
A smaller memory footprint also reduces memory access energy, critical for battery powered IoT devices [45, 11, 31, 5].
Fig.
1 shows
Pruning
UDC : U nified D NAS for C ompressible Models
Search Algorithm
Training Algorithm
Zero Weight Compression (Run-Length + Golomb-Rice)
Non-Zero Weight Compression (Golomb-Rice)
Compiler (Vela)
TFLite
QuantizedPruned
TFLite
Compressed
Compressed Model (Flash Memory)
Arm U55 NPU
De-Compression
NPU
Quantization 1.
Design Compressible Model 2.
Compress Model Weights 3.
Deploy on Device
[Figure]
HW
[Figure]
Operators
[Figure]
D e p l o y e d M o d e l
S i z
e ( M
B )
[Figure]
how model size scales with quantization, pruning, and both, when compressing a MBNetV [35] for deployment using Vela, down to a 9 . 3 × reduction in model size for 1-bit quantization and 1% non-zero weights in the extreme.
To fully exploit NPU model compression, we define our problem statement as follows:
[Equation]
the smallest model size
[Equation]
Typically, the model architecture itself is developed first, either through manual trial-anderror or by neural architecture search (NAS) [47, 21, 68, 33, 54, 67, 26, 62, 65].
Then, model conditioning techniques, e.g., quantization [61, 69, 64, 66, 63] and unstructured pruning [18, 43, 32, 31, 50], are applied before deploying to the target HW.
Furthermore, the generated model must also be small enough to deploy on the NPU device; merely regularizing size does not guarantee that the models found will fit into device Flash memory.
It is essential that:
Generated NNs do not exceed a hard constraint on (compressed) model size.
(P2)
This paper describes UDC (Fig. 1), which merges the model design and conditioning steps by conducting a joint search over NN architecture, weight bitwidths, and sparsity rates.
UDC builds upon differentiable NAS (DNAS), which exploits efficient weight sharing to solve (P1), while addressing key challenges like how to explore the design space while still addressing (P2) [51, 56].
The search space is HW-compression-aware, including only model conditioning techniques supported by the real-world Vela NPU compiler.
We explicitly exclude low-rank matrix factorization [38] and non-uniform floating-point quantization [66, 63], which are not supported.
The contributions of this paper are further summarized below.
Joint network architecture and conditioning search We extend the DNAS formalism to learn layer-wise weight sparsity levels.
We present a method for searching for sparsity levels in conjunction with layer bitwidths.
We show how to maximize weight sharing while jointly searching over sparsity, bitwidth, and layer width, as well as provide a differentiable and easily computable measure of compressibility for the DNAS algorithm to optimize.
Novel search algorithm
Previous work (e.g. [22, 21]) fails to effectively trade-off accuracy with model size (P2) in our search space (Table 3, first row).
UDC addresses this, with the following improvements: 1) guarantees that the search yields a model which satisfies specified
HW constraints (Sec. 4), 2) provides control over exploration-exploitation (Sec. 4.1), and 3) avoids over-regularization from biased Gumbel-softmax approximation (Sec. 4, 4.2).
Novel sparse, low-bitwidth representation and training algorithm UDC yields com- pressible models that are difficult to train.
We identify the root cause of the problem and propose a solution using a novel weight representation (Sec. 5).
SOTA NPU-deployable TinyML models Using UDC , we demonstrate compressible NNs Pareto dominant over prior work on CIFAR100, ImageNet, and super resolution (SR) tasks (Fig. 4)
Compression gains are validated by deploying to Ethos-U55 NPU.
[Table]
Table 1 contrasts UDC with relevant previous work.
MCUnet [48] is the closest to UDC , as it targets the TinyML HW form-factor with severely limited memory and includes results on large scale datasets like ImageNet.
However, UDC is fundamentally different to MCUNet: 1) unlike MCUnet, UDC also performs perlayer bitwidths and sparsity rate search to exploit HW compression, 2) UDC is a DNAS algorithm whereas MCUnet uses evolutionary search of a pre-trained once-for-all (OFA) supernet [17].
In particular, the UDC search space targets compressible models and is much more diverse than that of MCUNet; training a single OFA supernet for our search space would be infeasible (Sec. 6.2).
Yang et al.
[66] also demonstrate learning of per-layer bitwidth and sparsity rate, but otherwise differ significantly in that they: 1) do not include storage of the (required) pruning mask in the reported model size, which can even dominate at high sparsity rates [16]; 2) only consider model conditioning, without accompanying NN width, depth, or operator search; 3) employ floating-point non-uniform quantization, not deployable on TinyML MCU and NPU platforms which only support integer operations.
APQ proposes a multi-stage algorithm to search over layer width and bitwidth [64].
Compared to APQ, UDC : 1) optimizes over sparsity rates whereas APQ does not, and 2) produces models 3–9 × smaller than APQ.
Other relevant works include Gong et al.
[27], which searches over width and bitwidth but not sparsity, Choi et al.
[18], which searches over bitwidth and sparsity but not layer widths and without addressing (P2), and Uhlich et al.
[61] which searches over only bitwidths.
Modelling design decisions
We use bold-face to denote vectors/tensors, and z [ k ] to mean the k ’th element of vector z .
For a layer with input x and parameters θ , we denote its output by f ( x , θ ).
We say z ∈ R K follows a categorical distribution parameterized by π , i.e. z ∼ Cat ( π ), if p ( z [ k ] = 1) =
π [ k ].
Width selection Layer widths are modelled as f ( x , θ ) ⊙ w , ∥ w ∥ 0 / | w | =
ρ , where w is a binary mask to toggle any given channel, ρ is the fraction of non-zero channels, ⊙ is elementwise multiplication, ∥ w ∥ 0 is the ℓ 0 pseudo-norm which counts the number of non-zeros in w , and | w | is the number of elements in w .
We adopt the convention of setting the first ρ fraction of the channels of w to 1 [62].
Different choices { ρ 1 , · · · , ρ K w } of layer width can be modelled by the random variable (RV)
[Equation]
We refer to z w as a decision variable and to ρ k as an option.
Sparsity Sparse weight tensors are expressed as θ ⊙ m , ∥ m ∥ 0 / | m | =
s where m is a binary mask and s is the fraction of non-zeros.
We set the non-zeros of m to correspond to the largest magnitude elements of θ [50, 31, 28].
We model different choices of s amongst { s 1 , · · · , s K s } with the RV
[Equation]
Quantization Uniformly quantized tensors are given by
Q
[Equation]
, b,
[Equation]
where r is the quantization range and b the bitwidth.
We model different choices of b ∈ (cid:8) b 1 , · · · , b K q (cid:9) as the RV q ( π q ) =
P K q k =1 z q [ k ] Q ( θ , b k , r k ) , z q ∼ Cat ( π q ) , where each bitwidth b k is parameterized by its own range r k .
Sparsity and quantization Sparse, quantized tensors are modeled as Q ( θ , b, r ) ⊙ m .
We find setting the non-zero values of m based on θ superior to using Q ( θ , b, r ), which discards information about the relative magnitude of weights during the quantization process.
To model tensors whose sparsity level and bitwidth must be chosen, we define the RV q ( π q ) ⊙ m ( π s ).
[Figure]
Operator selection We frame the choice over different layer operators as the RV
P K f k =
z
f
[ k ] f k ( x , θ k ) , z f ∼ Cat ( π f ) , where each operator has a (possibly) different functional form f k ( · ) parameterized by its own weights θ k .
By including identity as a candidate operation, we can also model varying
NN depth.
Fig.
2 shows the layer-level search space, excluding operator selection for brevity.
Weight sharing is maximized by using the same θ in the sparsity, bitwidth, and width decisions.
The output of the layer in Fig.
2 is modelled as the RV f ( x , q ( π q ) ⊙ m ( π s ) ) ⊙ w ( π w ) .
3. Computing layer storage size
Satisfying (P1)-(P2) necessitates a quantitative measure of storage size, accounting for data compression.
For a given layer, the storage size achievable by prefix-free compression is lower-bounded by the weight (empirical) entropy, H ( · ), times the number of weight elements [19].
The entropy bound can, in turn, be bounded by (Appendix C):
[Equation]
The advantage of the right hand side (RHS) of (4) over the left hand side (LHS) is that it can be computed cheaply without processing θ .
Using the LHS requires computing the empirical entropy, which is expensive for large θ .
Moreover, combining the LHS with a gradient-based optimizer requires gradient approximation, since the empirical entropy is not differentiable [60].
We refer to the RHS of (4) as the compressed tensor size and use it as our measure of layer size.
Sec.
6.2 confirms that it is achievable with both an arithmetic encoder and Vela.
Let ϵ ( s, b, ρ ) be a given layer’s storage size (the RHS of (4)) as a function of sparsity s , bitwidth b , and non-zero channel fraction ρ .
When ( s, b, ρ ) must be chosen, let ϵ (cid:16)P K s k z s [ k ] s k , P K
q k z q [ k ] b k , P K w k z w [ k ] ρ k (cid:17) be the storage size.
For the entire NN, we sum the storage size of all layers and denote the result E ( { z } ), where { z } is short-hand for the set of all decision variables.
Proposed DNAS algorithm
Our optimization objective is:
[Equation]
where L ( · ) is a task loss, e.g. cross-entropy, D is the training data, and e ∗ is the target model size.
The constraint in (5) stems from the two-stage process typical in DNAS: 1)
Optimize { π , θ } and extract the most likely configuration { γ ( π ) } , 2) Train and deploy the result [51, 21, 22, 40, 62, 65].
While prior works use the constraint [21, 22, 40, 62, 27]
[Equation]
we constrain the most likely configuration since this is what is actually deployed.
Even if (6) is satisfied, it is not guaranteed that { γ ( π ) } , the deployed model, satisfies the constraint.
Gradient-based optimization Solving (5) using stochastic gradient descent (SGD) requires: 1) ensuring that the constraint is met, 2) differentiating with respect to { θ , π } .
To deal with the constraint and avoid dealing with the non-differentiable γ ( π ), we modify (5) to arg min
[Equation]
While L z E has a different form than the constraint in (5), it actually represents a much stronger constraint and its minimization implies the constraint is met.
Lemma 4.1.
If L z E = 0 , then for any sample of { z } , denoted { z s } , with non-zero probability,
E ( { z s } ) =
e ∗ and E ( { γ ( π ) } )
=
e ∗ .
L z E differs from (6) because it penalizes any configuration which violates the constraint, whereas (6) penalizes { π } only if the expected storage size violates the constraint.
The derivative of (7) with respect to (w.r.t.)
{ θ } can be approximated using a Monte-Carlo (MC) approximation of the expectation and applying standard automatic differentiation.
The derivative w.r.t.
{ π } is more complex since the expectations in (7) depend on { π } .
A popular solution is to use a biased but differentiable approximation of z , given by the Gumbel-softmax distribution: ˆ z = softmax (cid:16) log π
+ g τ (cid:17) , g [ k ] ∼ Gumbel (0 , 1) , τ > 0 [39, 21, 22].
As τ → 0, ˆ z approaches z in distribution, while the variance of any gradient estimator which uses ˆ z increases ([55] and Fig. 5d).
As a result, the common practice is to anneal τ from a high to a low value throughout the search.
Using ˆ z can lead to co-adaptation of search space options, which is undesirable but solvable in practice (Appendix D).
[Table]
Gumbel-softmax and over-
regularization We observe two issues with annealing τ : 1) increased gradient variance at low τ , coupled with a complex search space, causes issues for SGD, 2) when we replace z with ˆ z in L z E , with the result
L ˆ z E , the regularizer becomes artificially inflated.
To understand the impact of τ on L ˆ z E , we evaluate L ˆ z E for different values of τ , setting π such that E ( { γ ( π ) } )
=
e ∗ .
The results are presented in Table 2, col. 1 and show that increasing τ increases L ˆ z E , i.e. the relative impact of L ˆ z E on (7) depends on τ .
As such, we seek to keep τ low, while minimizing gradient variance.
Our solution is to use multiple samples of { ˆ z } in the MC approximation of (7).
To maintain the same computational cost as the single MC sample case, we divide the number of optimization steps by the number of samples.
As well as reducing gradient variance, our strategy has two additional practical benefits: 1) trivial extension to multi-GPU systems, since each GPU can run its own MC sample and gradient computation, 2) the overheads of computing the gradient are amortized across the MC samples.
We observe a 5 .
3 × speed-up when going from 1 MC sample to 32 (Table 3, col. 6, row 4 vs. row 8.)
4. Exploration-exploitation
When solving (7), the goal is to explore as many configurations as possible (exploration), while still training each configuration for a meaningful number of steps (exploitation).
Ideally, the search algorithm should gradually move from exploration to exploitation.
We propose to explicitly control the exploration-exploitation trade-off by projecting π onto the set
[Equation]
(8) after each SGD step, where ξ t is the upperbound on π at step t .
Setting ξ t = 0 constrains π to parameterize a uniform distribution and represents maximal exploration.
Setting ξ t = 1 − 1 / | π | removes the constraint on π , allowing the optimizer to enter full exploitation.
We define the projection operator as P S ( π ) = softmax (log π /T ∗ ),
T ∗ = arg min T P | π | k max ( P S ( π )[ k ] − (1 / | π | + ξ t ) , 0), which we solve numerically.
We choose this form for P S ( π ) because of its simplicity and because the relative ordering of options between π and P S ( π ) does not change.
While we find it necessary to enforce exploration by projecting π onto S , L z E implicitly promotes exploitation.
Lemma 4.2.
Let L z E =
and z j be the j ’th decision variable.
Let there be no decision for which two of its options have the same cost, i.e. for two configurations { z s } and n z s ′ o
such that z sj ̸ =
z s ′
j and z sk =
z s ′ k ∀ k ̸
= j , we have L z s E ̸
=
L z s ′ E .
Then each π j must be one-hot.
The assumption in Lemma 4.2 that no decision has two options with the same storage cost is satisfied for a typical compressible model search space.
4. Combating over-regularization with rejection sampling
Projecting π onto S enables explicit control over the exploration-exploitation dynamics, but also inflates L ˆ z E .
Table 2, col. 2 shows that setting ξ t = 0 .
5 increase L ˆ z E dramatically.
By forcing π to be closer to uniform, the number of configurations { ˆ z } with non-zero probability increases, such that the probability of a randomly drawn configuration violating the constraint also increases.
Inflating L ˆ z E forces the optimizer to focus less on L task , leading to solutions which meet the constraint but perform poorly on the target task.
L ˆ z E increases when the properties of { π } change because it depends on all possible configurations { ˆ z } instead of the most likely one, i.e. { γ ( π ) } .
To motivate the remedy, observe that while not all samples of ˆ z , denoted ˆ z s , satisfy γ (ˆ z s ) =
γ ( π ), some do.
Indeed, L ˆ z E would be 0 if it was evaluated over those samples that satisfy γ (ˆ z s ) =
γ ( π ), assuming E ( { γ ( π ) } )
=
e ∗ (and κ = 1 in Appendix D (12)).
We refer to samples generated in this manner as ˜ z
s
and they correspond to a RV whose distribution is different from ˆ z , but still depends on π and can therefore be used to generate gradients to π from L ˜
z E .
Alg.
1 shows how to generate ˜ z s .
Replacing ˆ z with ˜ z for all decisions negates the effects of controlling π through projection onto S .
Therefore, we use ˜ z for a given decision with probability
ϑ and ˆ z otherwise.
Algorithm
Sampling 1: Sample
ˆ z s ∼ p
(ˆ z ), 1 ≤ s ≤ S 2: ˜ s = 1, k ∗ = arg max k π
[ k ] 3: for 1 ≤ s ≤ S do 4: ˆ k = arg max k
ˆ z s [ k ] 5: if ˆ k = k ∗ then 6: ˜
z
˜ s
= ˆ z s , ˜ s =
˜ s +7: z = 1˜ s − 1 P ˜ s − 1 s ′ =
1 ˜ z s ′
Average 8: return ˜
z
Algorithm
1: for 1 ≤ t ≤ t max do 2: for 1 ≤ s ≤ S do 3:
Generate MC sample s for decision j using ˆ z j w.p. ϑ and ˜ z
j else 4:
Update { m } if t modulo 16 =5:
Take SGD step on (7) 6: π ← P S ( π ) ▷ Projection 7: return { π }
Using ˜ z has two major benefits.
First, even when exploration is enforced by projecting π onto S , L ˜ z
E is not inflated.
Table 2, col.s 3-4 show how ϑ > 0 brings down L ˆ z E for the same underlying { π } .
Second, mixing ˆ z and ˜ z gives the flexibility of being in exploration for some decisions and exploitation for others.
By randomly choosing which decisions use ˆ z and which use ˜ z , we prevent greedy behavior whereby a given decision enters exploitation and never returns to exploration.
The complete UDC algorithm is summarized in Alg. 2.
Training sparse, quantized models
DNAS is typically a two-stage process: a search to find the model architecture is followed by finetuning to find optimal weights.
However, we find that the second stage yields poor results when training sparse, quantized models from scratch [71].
UDC employs a threestage finetuning process: Stage 1: initialize { θ } and train with quantization enabled but unstructured pruning disabled, Stage 2: enable unstructured pruning gradually, Stage 3: train with both quantization and unstructured pruning enabled.
To counter the training challenges induced by quantizing and pruning θ , we employ several known techniques, with one slight modification.
For quantization, we use a variant of [59] where each weight is quantized with probability α during the forward pass, such that the weights used for training are Q ( θ , b, r ) ⊙ h + clip ( θ , − r, r ) ⊙ ( 1 − h ) , h
[ i, j, c in ,
c out ] ∼ Bernoulli ( α ).
Unlike UDC , [59] uses θ in the second term.
We observe that the range of θ can differ significantly from that of Q ( θ , b, r ), especially when unstructured pruning is applied during the learning process, so we clip θ to the same range as Q ( θ , b, r ).
For unstructured pruning, we gradually anneal the pruning rate from 0% to 100% of the target rate during stage 2 [72, 50].
Weight numerical representation To understand the challenge of training sparse, quantized NNs, consider a single layer of pretrained weights, with histogram in Fig. 3a.
The quantization bins ( b = 4) are shown in red and the pruning boundary in purple, with everything between the purple lines mapped to 0.
The sparse quantization problem is clearly apparent here: 8 / 14 of the non-zero quantization bins are unused because they fall inside the pruning boundary.
If we now train with weight sparsity and quantization constraints, the weight distribution, quantization bins, and pruning zone adjust (Fig. 3b).
The optimizer flattens the weight distribution to use more quantization bins, resulting in increased weight range.
To quantify the weight growth, we report the norm of the NN weights before and after training in Fig.
3, showing an increase of over 4 . 6 × .
The rate at which NNs can be trained, known as the effective learning rate, is inversely proportional to the weight norm [12].
Therefore, the interaction of sparsity and quantization cause weight norm inflation, which decreases the effective learning rate, reducing NN performance.
We propose a different weight representation and only quantize the range beyond the pruning boundary (Fig. 3c), using the quantization operator ˆ Q ( θ , b, r, β ) =
Q ( θ − sign( θ ) β, b, r ) + sign( θ ) β for b > 1 , β ∈ R .
We set β to be the largest pruned value of θ .
Training the sparse, quantized NN with the proposed number representation leads to much smaller weight norm (only 2 .
2 × growth over the pretrained weights), which makes training easier and accuracy higher (Fig. 5c).
Deployment with integer math NPU/MCU HW platforms typically only support integer operations, which are cheaper than floating point.
There are at least two ways of deploying
NNs quantized using the proposed approach on such HW.
Firstly, convolution can be decomposed into f ( x , Q ( θ − sign( θ ) β, b, r ))
+ βf ( x , sign( θ )).
Both terms can be calculated using only integers, but since sign( θ ) is a 1-bit tensor, the second term does not require any multiplications.
Secondly, the NN can be trained using the proposed number format, then deployed with Q (cid:16)
ˆ Q ( θ , b, r, β ) , b ∗ , r ′ (cid:17) , such that the deployed NN is uniformly quantized, with b ∗ = 8 to match the datatype supported in MCUs/NPUs [7, 8, 2].
The advantage of the latter approach is that the training benefits from the expressivity of the proposed number format, while its deployment uses a standard data type.
We find that this approach with b ∗ = 8 does not incur an accuracy loss on ImageNet (Fig. 4a).
Results
We compare UDC with SOTA methods on a model size vs. accuracy basis.
All reported model sizes, including related works, use compressed size (RHS of (4)), except Choi et al.
[18], who use bzip2 to compress weights and we use their reported sizes.
CIFAR100 is an image classification task with 50k training / 10k test images, andclasses.
Our search space is based on the wide residual network with depth 20 and width multiplier 10 [70, 52].
We search over layer width (increments of 10% of the orignal), bitwidth (1, 4, 8, 32), and sparsity (1 − 100% non-zeros, increments of 10%).
We use UDC ( settings in
Appendix E) to find and train models at two sizes: 0 .
55MB and 0 .
7MB.
Fig.
4b shows that
UDC generates Pareto-dominant models compared with SOTA methods [27, 52, 65, 63, 58].
ImageNet is an image classification task with 1 .
M training / 50k test images, and 1k classes.
Our search space is based on MBNetV2, with the same options as for CIFAR100, and we target 0 .
5, 1, and 1 .
25 MB models (Appendix F).
For the 1 . 25MB experiment, we replace 3 × 3 kernels with 5 × 5 and make all layers 1 .
5 × wider in the baseline architecture.
In practice, the model size constraints are determined by the Flash memory size of the deployment HW platform.
To be sure, 0 .
5 − 1 .
25 MB Flash sizes are fairly common for
[Figure] [Figure]
commodity HW platforms [6] and are often used in research targeting deployment on constrained HW platforms [48, 14].
As such, we targeted this range because it represents a reasonable, but extremely challenging deployment scenario.
Fig.
4a shows that UDC generates pareto-dominant NNs vs. the SOTA [48, 18, 27, 49, 65, 62, 64].
Next, we include depth/operators in the search space, considering two alternatives to the inverted bottleneck: 1) regular 3 × 3 convolution, 2) identity.
We target a 0.5MB model.
UDC chooses inverted bottleneck blocks everywhere, such that the model found by including depth/operators in the search is identical to the one found when they are excluded.
We compare UDC with non-uniform quantization approaches (Appendix H), and find that
UDC models are pareto-dominant even though approaches like HAQ [63] and [66] employ a more expressive, non-uniform quantization which cannot be deployed on commodity MCUs and NPUs with integer math operations.
Finally, comparing UDC to a SOTA unstructured pruning algorithm (Appendix I) shows that UDC finds considerably more accurate models.
SR The purpose of the SR experiment is to: 1) show that UDC can be applied to regression problems, 2) demonstrate UDC in a setup which constrains the computational complexity of the NN, measured in number of multiply-and-accumulate (MAC) operations, instead of model size.
Note that 1 MAC = 2 floating point operations (FLOPs), but authors often conflate the two terms.
This experiment excludes bitwidth and sparsity from the search space, which do not impact the inference compute cost.
The topology of our search space is inspired by FSRCNN [20], with search over depth, width, and kernel size.
We do all search/training on div2k and report results on set14 (see Appendix G for results on div2k and set5).
We target 4x upscaling and report MACs for an input patch of 64x64.
We train all
NNs using the same settings.
Fig.
4c shows that UDC finds Pareto-dominant NNs compared with SOTA efficient SR methods SESR, FSRCNN, and TPSR [15, 20, 44].
6. Comparison with random search
Random search is a strong baseline that can compete with DNAS in some settings [46].
We are not able to add early stopping [46], since sparsity is only applied during stages 2-3, when it is too late to save time from stopping.
Fig.
5a compares UDC to randomly generated
[Figure]
[Table] [Table]
[Table] [Table]
NNs constrained to model sizes < 0 .
55MB for CIFAR100, with the search space used in Fig.
4b.
The gap between the best UDC models and random search is 4 .
The best accuracy achieved by random search after a given number of trials is shown in Fig 5b.
To integrate out randomness from the the ordering of trials, we permute the order of the random search results 100 × , averaging over the permutations.
Since UDC search requires ∼ 2 × longer to find a NN than training a baseline NN, we plot the cost of UDC as 3 trials (2 to search and 1 to finetune).
These results confirm that UDC finds better NNs faster than random search.
DNAS alternatives like evolutionary search require training many NNs or proxies to relax the compute burden [24, 23, 17, 48, 49].
However, using proxies assumes a search space dense with performant NNs, which is not true in our case (Fig. 5a).
6. Validation of algorithm and system components
[Table]
Compression We verified our model size approximation, i
[Table]
e. the
RHS of (4) with pruning mask compressed to the entropy limit, by compressing m using arithmetic coding [53] for all UDC ImageNet models (Fig. 4a).
We observed that theoretical and practical compressed model sizes are within two decimal places (Table 4).
Tableshows that Vela can achieve compression ratios even larger than those predicted by UDC for NPU deployment, which stems from the fact that non-zero weights can often be compressed to less than b bits per weight by Golomb-Rice compression.
Ablations Table 3 gives an ablation on UDC components (Sec. 4) on CIFAR100 and ImageNet, which we evaluate on final accuracy, hyperparameter search over λ requirement, and runtime in GPU days (GPUD) normalized to a typical training setup.
Table 3 shows that disabling
[Table]
any of the UDC components leads to significantly worse results on both datasets.
Vanilla
DNAS [21] fails to meet the HW constraint.
Fig.
5c shows the isolated benefit of the proposed number format over the baseline (3).
Runtime An important aspect in evaluating NAS algorithms is their runtime, or the time required to yield results.
Part of the challenge in comparing algorithm runtimes is that the runtime generally depends on: 1) software implementation quality and HW platform, which jointly form the “system” that the algorithm runs on, 2) the algorithm itself,
i
[Table]
e. what the algorithm is actually doing during the search, 3) the number of epochs (or amount of data processed) for which the search is run.
The challenge is that only 2-3 are algorithm dependent, whereas 1) depends on the code quality and the resources of the experimenter (i
[Table]
e. higher grade GPUs exhibit higher throughput compared to low-end GPUs).
In order to compare with other works
[Table]
we now attempt to disentangle the 3 components that make up the runtime in Table 5.
We list the system and algorithm specific search speed measured in images per second, when available from the reference, the number of search epochs, the search cost in GPUD under the system used in the reference, and the search cost in GPUD under a common system assumed to be running at 460 images / s [4].
For several references, the algorithm begins with a pretrained model, which we assume is trained for a standard 200 epochs [4].
Table 5 shows that in absolute terms (i.e. the system specific cost), UDC is 1 . 4 × faster than FBNet and FBNetV2.
When comparing approaches based on a normalized system running at 460 images / s, UDC is faster than all of the competing approaches other than FBNet and FBNetV2.
Our hope is that Table 5 gives a rough sense of the relative search cost of UDC and the competing methods.
Societal impact & limitations Developing UDC used hundreds of energy-consuming GPU hours.
However, this can be amortized by increasing the energy efficiency of billions of
IoT devices.
A limitation of UDC is that we retrain { θ } for every constraint e ∗ , whereas approaches based on OFA can amortize the cost of training { θ } across multiple constraints.
Conclusion
Emerging NPU HW platforms specialized for TinyML support model compression, whereby quantized and pruned NNs can be stored in a reduced memory footprint.
While compression is highly desirable, it increases the complexity of the NN design process, as the space of candidate NNs is increased by adding quantization and pruning on top of the conventional NN architecture choices.
To enable TinyML practitioners to fully exploit HW model compression in NPUs, we describe a unified DNAS framework to search both architecture choices and aggressive per-layer quantization and pruning.
We describe a number of improvements on top of DNAS, allowing us to demonstrate SOTA TinyML models that fully exploit model compression, as well as a comparison with random sampling and extensive ablations.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] See Sec. 6.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See Sec. 6.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix A and B 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] The code is proprietary.
(b) Did you specify all the training details (e
g., data splits, hyperparameters, how they were chosen)?
[Yes] See Sec.
6 and Appendices E-G
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] The computational cost of running NAS repeatedly to generate error-bars exceeds our computational capabilities.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] We report the search cost, normalized to a typical GPU NN training setup of 460 images / second, in Table 3.
We also report the search cost in absolute terms in Table 5. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] We used the CIFAR100, ImageNet, div2k, set5, and set14 datasets with the appropriate creators cited.
(b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Proof of Lemma 4.
Proof.
We know that p ( E ( { z } ) ̸ = e ∗ ) must be 0.
But if E ( { γ ( π ) } ) ̸ = e ∗ , p ( E ( { z } ) ̸ = e ∗ ) > 0, which is a contradiction.
B Proof of Lemma 4.
Proof.
If there is a π j which is not one-hot, then the following configuration sample has non-zero probability:
• z sj ̸
= s s ′
j • z sk =
z s ′ k ∀ k ̸
=
j
Since we assumed that L E ( { z } ) = 0, Lemma 4.1 gives that E ( { z s } ) =
e ∗ =
E (cid:16)n z s ′ o(cid:17) .
But this is a contradiction since we assumed there are no decisions for which two options have the same efficiency.
C Derivation of entropy bound
The entropy bound in (4) can be derived as follows:
[Equation]
where the first inequality follows from the fact that the entropy of a product of RVs is bounded by the sum of their entropies and the second bound follows from the fact that
Q ( θ , b, r ) costs at most b bits per element to encode.
D Avoiding co-adaptation in DNAS
While samples of z are one-hot, samples of ˆ z are not.
This property can cause issues for an approach like DNAS, where weight-sharing can lead to co-adapation between search options [34, 29].
The result is a large performance drop when finetuning { γ ( π ) } , compared to the value of L task achieved by solving (7) [36].
One solution, which we adopt, is to use a straight-through-estimator (STE), whereby
[Equation]
else
[Equation]
is used in the forward pass, where κ is the number of non-zeros in ˆ z f
[ k ], topk(ˆ z , κ ) returns the indices of the κ largest elements of ˆ z , and ˆ z is used in the backward pass [39].
Typically, κ ∈ { 1 , 2 } .
E CIFAR100 experiment settings
We run the search for 200 epochs, annealing τ from 0 .
66 to 0 .
1 using an exponential schedule.
We use SGD for θ with learning rate annealed from 0 .
1 to 1 e − 4 using a cosine schedule
and we use ADAM for { π } with a constant learning rate of 1 e − 3.
We increase ϑ from 0 to 0 .using a linear schedule
and we increase ζ t from 0 .
1 to 1 using a linear schedule.
We initialize the search by running a warmup stage for 50 epochs where we use SGD for θ with learning rate annealed from 0 .
1 to 1 e − 4 using a cosine schedule, τ is annealed from 0 .
66 to 0 .
1 using an exponential schedule, { π } is not learned, and ϑ = 0.
During both warmup and search, we set κ = K for all width decisions and κ = 2 for all quantization and sparsity decisions.
[Table]
To finetune the discovered models, we run stage 1 for 254 epochs using SGD and cosine decay with restarts learning rate schedule, cycling between 0 . 1 and 1 e − 4 at intervals which double after every cycle and beginning with a cycle of 2 epochs.
We run stage 2 for 60 epochs, annealing the learning rate from 0 .
1 to 1 e − 4 using a cosine schedule
and then we run stage 3 for 30 epochs, annealing the learning rate from 1 e − 4 to 0 using a cosine schedule.
We use distillation with a teacher model whose architecture is WRN 20-10.
For data augmentation, we use horizontal flipping, random crop with a size of 4, and cutout with a size of 18.
We use ℓ 2 regularization with coefficient 5 e − 4.
We disable learning of batchnorm scale and offset parameters [52].
Table 6 presents the detailed experimental results for the CIFAR100 experiments.
F ImageNet experiment settings
We use the same search settings as for the CIFAR100 experiments.
To finetune the discovered models, we run stage 1 for 126 epochs using SGD and cosine decay with restarts learning rate schedule, cycling between 0 . 1 and 1 e − 4 at intervals which double after every cycle and beginning with a cycle of 2 epochs.
We run stage 2 for 60 epochs, annealing the learning rate from 0 .
1 to 1 e − 4 using a cosine schedule
and then we run stage 3 for 30 epochs, annealing the learning rate from 1 e − 4 to 0 using a cosine schedule.
For the 0 . 5 and 1MB target experiments, we use distillation with a teacher model whose architecture is MobileNetV2.
We do not use distillation for the 1 . 25MB target experiment.
For data augmentation, we use the standard ImageNet data pipeline [22], as well as horizontal flipping and label smoothing with smoothing coefficient 0 .
We use ℓ 2 regularization with a coefficient of 1 e − 4.
Table 7 shows the detailed ImageNet results.
G Super resolution experiment settings
We run the search for 300 epochs, with constant τ set to 0 .
We use ADAM for θ with learning rate annealed from 1 .e − 4 to 1 e − 5 using a cosine schedule and ADAM as well for { π } with a constant learning rate of 1 e − 3.
We keep ϑ constant to 0 .
and we increase ζ t from 0 .
1 to 1 using a cosine schedule.
Table 8 provides detailed results for the super resolution experiment.
H Comparison to non-uniform quantization approaches
We compare UDC to approaches which employ non-uniform quantization in Fig.
UDC is Pareto-dominant even though it uses uniform quantization and can be deployed on
MCUs/NPUs with integer math whereas the other approaches cannot.
[Table]
[Figure]
[Table]
I Comparison to approaches which only do unstructured pruning
We also compare UDC to a SOTA unstructured pruning algorithm [43] in Fig.
As the results show, UDC generates much more accurate models.
J Visualization of design choices
We provide a visualization of the design choices made by UDC for the ImageNet experiments in Fig.
[Figure]
[Figure] [Figure]