Inspiration
I've spent the last while building deep learning infrastructure from scratch — CUDA kernels, autograd systems, memory pools — because I wanted to actually understand what happens under the hood of the frameworks everyone uses. When the Arm AI Optimization Challenge came up, it felt like the natural next step: take that low-level systems instinct and point it at a real, unglamorous engineering problem — models are big, edge devices aren't, and Arm CPUs are increasingly what's actually running inference in the real world, from phones to Arm-based cloud instances. Low-rank factorization stood out because it's mathematically clean (SVD is something I already understood deeply) but has a genuinely uncertain payoff on real hardware — which made it an interesting thing to measure, not just implement.
What it does
TensorFold-arm is a tool that works with PyTorch models. It makes these models smaller by changing the layers. It does this with a method called SVD. This method replaces a weight matrix W with two smaller matrices A and B. The weight matrix W is similar to the product of A. B.
The tool does not change every layer. It only changes a layer if it actually makes the model smaller. Some layers are not worth changing. The tool checks this automatically. It does not just assume that every layer should be changed.
We tested TensorFold-arm on a model. This model is a 3-layer MNIST MLP. The results were:
The model had 535,818 parameters. After using TensorFold-arm it had 279,940 parameters. This is a reduction of 47.75%.
The model was 97.79% accurate. After using TensorFold-arm it was 97.06% accurate. This is a decrease of 0.73%.
The model ran faster after using TensorFold-arm. It was up to 1.41 times faster on ARM64 hardware. This was tested with batch sizes from 1 to 32. This is how models are usually used on devices.
We got these results from a workflow, on GitHub Actions. This workflow runs on ubuntu-24.04-arm. This means it runs on Arm hardware, not a simulator.
How we built it
The core is a layer that changes a dense matmul for something like this: (X @ A) @ B plus bias. It automatically chooses the rank based on how much energy from the singular values needs to be kept which's 90 percent by default. The compress function looks at a model does this analysis on every nn.Linear and swaps in the factorized version when it really helps.
We made a benchmarking harness after we got the base compression pipeline working. This harness has threads pinned to 1, 100 warmup iterations, 500 timed iterations times 10 repeats. We connected it to ARM64 CI so anyone who uses the repo can get the results, not just something we ran once on our own computer and put in a README.
When we saw that TensorFold was getting slower than dense at batch sizes on ARM64 we did not just want to report that and stop. We thought it might be because of the overhead of sending work to the kernel since it has two matmul calls per layer of one. So we made a custom CPU kernel using C++ and pybind11. This kernel is made in time for each platform and we checked that it is correct on 5 different shapes and 5 batch sizes, on both Linux and Windows. We did this to test our idea directly.
Challenges we ran into
We had a big problems to solve. The first one was our fault. Our main README file and the detailed benchmark results file did not agree with each other. One said that the speedup increased steadily up to 1.664× at batch 256 but the other file, which had the numbers that we verified showed that TensorFold was actually slower than dense at that batch size. We caught this mistake by checking our documents against each other which saved us from looking bad in front of the judges.
The second problem was with the tools we used to build extensions on different platforms. On Linux we needed to use GCC-style flags like -fopenmp and -march=but on Windows we needed to use MSVC-style flags like /openmp. If we sent the flags to the compiler it would just fail. Also on Windows the MSVC compiler only works when we launch it from the "x64 Native Tools Command Prompt for VS". This cost us a lot of time to figure out before we found a way to build our project that worked on all platforms and gave us an error message instead of just failing.
The biggest problem we had though was not technical. More about understanding what was going on. After we built and tested the fused kernel we found out that it did not make any difference on the ARM64 benchmark. This was a thing to accept especially, during a hackathon. Our first thought was to keep trying to make it work until we got a result.. Instead we did a careful comparison of the dense original and fused kernels to really understand why the fused kernel was not helping. We did not want to release a kernel that did not do anything so we took the time to understand what was going on with the TensorFold and the fused kernel.
Accomplishments that we're proud of
We found a mistake in our benchmark before someone else did. We got a working CPU kernel that's correct and it works on two different compilers and platforms. What is also important is that we did an experiment got an answer we did not like and we wrote about it instead of hiding it. The CPU kernel and its 28 tests are still in the repository they work and we have documents, for them. We just do not say that they make things faster because the data does not show that the CPU kernel makes things faster.
What we learned
Parameter reduction and inference speedup are. They are not the same thing. A model that is 47.75% smaller does not automatically mean it will have inference at every batch size. If we had treated them as the thing our results would have been misleading. We also found out that removing kernel-dispatch overhead is not always a thing. Our fused kernel traded batched-GEMM efficiency for dispatches.. These two effects pretty much canceled each other out. The real problem at batch sizes is more basic. It is about how the FLOP savings from a given rank compare to how a single dense batched GEMM scales with batch size, on the hardware. This is not something you can fix by fusing things together. Parameter reduction and inference speedup are still important. We need to think about them separately.
What's next for Tensorfold
Broader model coverage. Now everything is validated on one MNIST MLP; extending to a small CNN or transformer block would test whether the compression and speed characteristics generalize.
Arm-optimized inference backends. Routing the model through ONNX Runtime with the XNNPACK execution provider to see whether an Arm-tuned GEMM implementation succeeds where our own fused kernel didn't.
Stacking int8 quantization on top of the low-rank factors since Arm cores with dot-product extensions get hardware acceleration quantization alone can't unlock from fp32 factorized weights.
Accuracy-aware rank selection. Moving beyond SVD-energy- rank choice toward measuring actual downstream accuracy impact, per layer.
Memory footprint measurement. Process RSS, not just parameter count since thats the number that matters most for the Mobile AI use case this project is scoped around.
Log in or sign up for Devpost to join the conversation.