Axon TFLite compiler: largest layer size unexpectedly doubled

[owadekimo]

Hi,

While using the Axon TFLite compiler to compile a series of TFLite models, we encountered a peculiar issue.
The largest layer in the models (i.e. the dense layer right after the convolutional part of the network, just after flattening) is consistently doubled in size with respect to its original TFLite model.
Closer inspection reveals that this specific layer indeed has double the number of original weights, but half of them are set to 0 in a consistent [x, x, 0, 0, x, x, 0, 0, ...] pattern (see the compiler output). This effect is persistent across all our 64 tested model sizes from ~4KB to ~1.2MB.

Because we don't see correctness of the model impacted by this, it seems like this is a deliberate action taken by the compiler.
It seems possible to us that this may be done as some kind of internal weight reordering operation in the Axon hardware.

We'd like to know if this doubling is a premeditated action by the compiler and if it can possibly be prevented in some way.
Attached are the compiler logs and the intermediate and final compiler output for the largest model we tried this on, which don't seem to log/warn about anything relevant to this phenomenon. 
Our thanks for any information you'd be able to give us about this issue.

Related issue

nrf_axon_model_model_test_uwb_binfile_content_.txt

Fullscreen run_axons_ml_nn_compiler_executor_20260623_181238.log Download

INFO: running uwb_embedded_test
INFO: starting compiler...
INFO: model name is : model_test_uwb
INFO: test data or proper test labels are not provided, we have the tflite file.
INFO: running for variant model_test_uwb
INFO: RESHAPE is a PASSTHROUGH Operator
INFO: compiler intermediate outputs are saved in : C:\dekimo\vlaio-uwb-edge-ai\test-nrf54-axon_npu\models\uwb_test_model_64/intermediate/
INFO: model compiled successfully ....
INFO: not running inference as test vectors are not provided
INFO: done running variant model_test_uwb, return code 0
INFO: 
	Performance_Metrics
		Model_Variant	Inference_Time	Accuracy	Precision	Quantization_Loss	Flash_Size	Ram_Size	Total_Memory	
		model_test_uwb	None	None	None	None	0	0	0
INFO: completed running uwb_embedded_test, return code 0(NRF_AXON_COMPILER_RESULT_SUCCESS), took 1.80 seconds

nrf_axon_model_model_test_uwb_.h

[tesc]

Hi,

Closer inspection reveals that this specific layer indeed has double the number of original weights, but half of them are set to 0 in a consistent [x, x, 0, 0, x, x, 0, 0, ...] pattern

This stems from a limitation that each output row must be on a 32 bit boundary. Since the output width of the previous node is 2, there is 2 bytes padding on each row. The following fully connected layer must account for this additional width in its flattened input. (Flattened length is 32 * 4, not 32 * 2.)

However, what you can do, and would likely benefit from, is to transpose the model. For the mentioned layer, the [32, 2] input would then become [2, 32]. The width 32 is a multiple of 4, so then there is no padding needed. Transposing the model will also improve the performance of other layers; Axon NPU prefers a wider aspect ratio vs a taller one. The compiler has a transpose option.

In order to transpose, you can add the following to the compiler .yaml file:
transpose_kernel: True

You then must also account for the transposing in the input and output data. (However in this particular case, where the output is one-dimensional, there should be no need to convert the output.)

Please note that the special case of width 1 is handled differently, without padding, so in that case the weights will not grow.

Regards,
Terje

[owadekimo]

Hi Terje,

Thank you for the quick and detailed response. I can confirm that turning on the `transpose_kernel` option both solved the issue and significantly increased the performance of the model. Good to know the detail about each output row being forced to 32 bit alignment through way of padding, this can help inform the way we build models for Axon.

Thanks again!