publications | Rohit K. Vartak

2025

Fed-SB: A Silver Bullet for Extreme Communication Efficiency and Performance in (Private) Federated LoRA Fine-Tuning

Raghav Singhal, Kaustubh Ponkshe, Rohit Vartak, and 2 more authors

2025

Abs Code

Low-Rank Adaptation (LoRA) has become ubiquitous for efficiently fine-tuning foundation models. However, federated fine-tuning using LoRA is challenging due to suboptimal updates arising from traditional federated averaging of individual adapters. Existing solutions either incur prohibitively high communication cost that scales linearly with the number of clients or suffer from performance degradation due to limited expressivity. We introduce Federated Silver Bullet (Fed-SB), a novel approach for federated fine-tuning of LLMs using LoRA-SB, a recently proposed low-rank adaptation method. LoRA-SB optimally aligns the optimization trajectory with the ideal low-rank full fine-tuning projection by learning a small square matrix (R) between adapters B and A, keeping other components fixed. Direct averaging of R guarantees exact updates, substantially reducing communication cost, which remains independent of the number of clients, and enables scalability. Fed-SB achieves state-of-the-art performance across commonsense reasoning, arithmetic reasoning, and language inference tasks while reducing communication costs by up to 230x. In private settings, Fed-SB further improves performance by (1) reducing trainable parameters, thereby lowering the noise required for differential privacy and (2) avoiding noise amplification introduced by other methods. Overall, Fed-SB establishes a new Pareto frontier in the tradeoff between communication and performance, offering an efficient and scalable solution for both private and non-private federated fine-tuning.

2023

Robustness to Variability and Asymmetry of In-Memory On-Chip Training

Rohit K. Vartak, Vivek Saraswat, and Udayan Ganguly

In Artificial Neural Networks and Machine Learning – ICANN 2023, 2023

Abs

In-memory on-chip learning is crucial for low-power, in-field training capabilities at the edge. We demonstrate the robustness of on-chip back-propagation to hardware variability in terms of bit-cell transistor }}V_T}}variability (}}2.5\backslashtimes }}more robust than off-chip training). We use perturbation schemes, asymmetry variations and variability-aware update schemes to identify the relative contribution of different on-chip operations: forward pass, backward pass and weight updates to Fashion-MNIST classification performance degradation with variations. It is revealed that variability during the weight update step is crucial while accuracy of backward pass or gradient calculation is not critical. We promote weight perturbation scheme over back-propagation as the choice for on-chip in-memory training with reduced points of failure and low cost of hardware.