A better way to resize AI updates for unequal devices

Small updates, different budgets

Federated learning lets clients train a shared AI model while keeping their training data local. Instead of repeatedly exchanging every model parameter, clients can use low-rank adaptation, or LoRA: they train two relatively small matrices whose product represents an update to a larger weight matrix. In conventional LoRA fine-tuning, the pretrained weights remain frozen.

The matrices share an internal dimension called the adapter rank. A smaller rank means fewer trainable parameters and fewer numbers to transmit, but also restricts the updates the adapter can represent. Clients with different budgets or data may therefore use different ranks.

That creates a server-side problem. After combining uploaded adapters, how should the server construct a smaller adapter for each client?

In an arXiv preprint, Shuaijun Chen and colleagues propose RBLA+, an extension of their earlier RBLA method. Its central change is how the server extracts client-sized adapters from its aggregate.

Why simply cutting can lose information

Original RBLA averages adapter positions using only clients whose adapters contain those positions. This avoids treating missing entries from smaller adapters as zero-valued observations. To serve a smaller client, it then keeps the first columns of one aggregated matrix and the corresponding rows of the other.

This is inexpensive, but the first positions are not necessarily the most useful ones. Different pairs of matrices can represent exactly the same complete update while producing different updates when shortened. The paper illustrates this with a small example in which changing the factor representation preserves the full update but increases the error of its rank-one prefix from 1 to 5.

Some error is unavoidable when a client’s rank is too small to represent the aggregate exactly. The authors distinguish that unavoidable loss from additional error introduced by the extraction method.

Their target is minimum Frobenius error: the square root of the sum of squared differences between entries in the original and reconstructed update matrices. This measures preservation of the update, not prediction accuracy.

Decomposing a small core

A standard mathematical tool, singular value decomposition, can find the best approximation at a chosen rank. But constructing and decomposing the full update matrix can be expensive.

RBLA+ instead uses an established compact procedure. Two QR factorizations express the adapter matrices using orthonormal bases and small coefficient matrices. The server combines those coefficients into a small core and applies singular value decomposition there, without forming the full update.

The resulting factors have components ordered by singular value. One decomposition therefore serves multiple client ranks: each client receives the leading components that fit its budget. In exact arithmetic, the full pair preserves the aggregate, and every shortened pair achieves the minimum Frobenius reconstruction error at its rank.

RBLA+ also adjusts both factors during upload and broadcast to account for different client and server scaling coefficients. Without this conversion, a receiving client could apply the right components at the wrong strength.

There is an important boundary to the guarantee. RBLA+ retains RBLA’s separate averaging of the two factors. Their product generally differs from an average of complete client updates. Spectral extraction optimally approximates that chosen aggregate; it does not fix the distinction between those aggregation rules.

What the experiments found

The researchers evaluated seven image datasets and six language tasks under two client-rank assignments. One gave larger adapters to clients with more samples; the other reversed that relationship. RBLA+ exceeded RBLA’s reported accuracy in 24 of the 26 settings.

For example, under the first assignment, peak CIFAR-10 accuracy rose from 38.33% with RBLA to 79.12% with RBLA+. However, the vision experiments trained low-rank factors from scratch, whereas the language experiments adapted pretrained TinyBERT. They are not all demonstrations of pretrained-model fine-tuning.

The tables report peak accuracy during training. Four MNIST-family datasets used three runs; the remaining tasks used one fixed seed, so small differences there do not establish reliable superiority.

Separate CPU benchmarks used synthetic adapters matching four model configurations. On an AMD Ryzen 9 9950X, RBLA+ took 3.5–11.6 milliseconds for extraction and preparation of ten client outputs. It was slower than direct RBLA truncation but faster than the tested decomposition alternatives.

Better reconstruction is not always better prediction

The paper directly demonstrates this limitation. Replacing RBLA’s broadcast with spectral extraction at round 100 on its MNLI training path reduced mean client accuracy from 56.81% to 42.60%. Repeated spectral broadcasting changes subsequent training, not just the current reconstruction.

RBLA+ preserves prescribed adapter sizes, so it adds no adapter payload over matched-rank RBLA. Its reported 90% downlink reduction relative to FLoRA applies to the evaluated ten-client transmission policy, not communication generally. The results support a server-side trade-off in simulations; deployed mobile-system costs remain future work.