Source-linked AI summary
Learning Private Neural Language Modeling with Attentive Aggregation
Shaoxiong Ji, Shirui Pan, Guodong Long, Xue Li, Jing Jiang, Zi Huang
TL;DR
Sensitive keyboard data motivates private neural language modeling without centralized collection, while simple federated averaging overlooks differences among client models. The paper introduces attention-based server aggregation that minimizes weighted server–client distances, and reports better perplexity and communication-cost performance than counterparts in most settings.
Problem
Centralized training raises privacy concerns for sensitive user data, while simple federated averaging ignores client importance and lacks server-side optimization for a well-generalized global model.
Method
The method uses attentive aggregation with layer-wise soft attention to weight client contributions and minimize weighted distances between server and client models.
Results
The proposed method outperforms FedSGD and FedAvg in testing perplexity across three datasets and outperforms counterparts in most comparison settings.
Takeaways & Limitations
Attentive federated optimization measures selected-client importance and can accelerate learning for private neural language modeling.
Abstract
from arXiv · showhide
Mobile keyboard suggestion is typically regarded as a word-level language modeling problem. Centralized machine learning technique requires massive user data collected to train on, which may impose privacy concerns for sensitive personal typing data of users. Federated learning (FL) provides a promising approach to learning private language modeling for intelligent personalized keyboard suggestion by training models in distributed clients rather than training in a central server. To obtain a global model for prediction, existing FL algorithms simply average the client models and ignore the importance of each client during model aggregation. Furthermore, there is no optimization for learning a well-generalized global model on the central server. To solve these problems, we propose a novel model aggregation with the attention mechanism considering the contribution of clients models to the global model, together with an optimization technique during server aggregation. Our proposed attentive aggregation method minimizes the weighted distance between the server model and client models through iterative parameters updating while attends the distance between the server model and client models. Through experiments on two popular language modeling datasets and a social media dataset, our proposed method outperforms its counterparts in terms of perplexity and communication cost in most settings of comparison.
I. INTRODUCTION
Mobile keyboard suggestion uses sensitive personal language data, making privacy-preserving federated learning attractive. The paper addresses limitations of simple model averaging with attention-based aggregation and server optimization.
- Centralized training collects user data on a server, creating privacy concerns if sensitive information is hacked or misused.
- Federated learning trains local models on client devices and learns a shared model without collecting data centrally.
- Mobile keyboard suggestion is a word-level language modeling task involving sensitive language preferences and private content.
- Simple federated averaging weights client models by sample counts despite differing language preferences and contributions.
- The proposed method uses attention to weight server–client model relations and minimize their expected parameter-space distance.
- Layer-wise soft attention captures the importance of parameters across local models during server optimization.
- Experiments report comparable perplexity and communication-round performance for private neural language modeling on mobile keyboard suggestions.
II. RELATED WORK
The related work spans federated learning, communication-efficient distributed training, neural language modeling, and attention mechanisms. The paper combines these areas for attentive federated optimization with GRU-based private language modeling and differential privacy.
- Federated Learning: Federated learning separates model training from data collection and supports privacy preservation through decentralized training.
- Federated Learning: Prior federated learning research includes client-side differential privacy, federated tensor factorization, and frameworks addressing distributed statistical challenges.
- Communication Efficiency: Communication efficiency research reduces federated learning costs through structured or sketched updates, quantization, and ternary gradients.
- Language Modeling: Neural language modeling research includes recurrent networks, LSTM models for long context, and weight tying to reduce trainable parameters.
- Attention Mechanism: Attention mechanisms guide perception and have been applied to recurrent vision models and sequence-to-sequence translation with global and local variants.
- This Paper: The paper applies attentive federated optimization to distributed clients, using GRU client models and a randomized mechanism for differential privacy.
A. Preliminaries of Federated Learning
Federated learning separates server aggregation from private local training. The framework sends a global model to selected clients, updates models locally using private data, and returns parameters for server aggregation.
- The federated framework contains server optimization and local training as its two main components.
- The server initializes model parameters, selects a client subset using the client fraction, and aggregates returned client parameters each round.
- Selected clients receive the server model, train GRU-based language models locally with stochastic gradient descent, and return model parameters.
- Local training keeps user data on client devices while model parameters are sent to the central server over a secure connection.
B. Attentive Federated Aggregation
Attentive Federated Aggregation introduces layer-wise attention to weight client-model contributions and iteratively optimize a global model close to selected client models in parameter space.
- Attentive aggregation: FedAtt introduces layer-wise attention for aggregating selected clients’ neural language models into the central global model.The mechanism operates on learned parameters of each neural-network layer.
- Server optimization: The optimization minimizes the weighted distance between the server model and client models using self-adaptive attentive weights.The server model parameters are updated by gradient descent over selected client models.
- Attention mechanism: The server treats its parameters as queries and client parameters as keys, deriving layer-wise attention scores from parameter similarity.Similarity is calculated from the norm of the difference between server and client parameter matrices, followed by softmax.
- Optimization procedure: FedAtt iteratively updates global server parameters from the current server model and client updates.Algorithm 3 returns updated global parameters after attentive optimization.
- Aggregation properties: The aggregation is fine-grained across neural layers, selecting client features according to similarity with the server model.This design is intended to produce a fine-tuned global server model that represents client models well.
C. GRU-based Client Model
The client language model uses a GRU-based recurrent architecture, with token embeddings processed through recurrent gates and hidden states.
- Client architecture: The client-side language model uses a gated recurrent unit (GRU), a simpler LSTM variant with merged gates and states.The learning process is model-agnostic, but this paper selects GRU for language modeling.
- Input representation: Input words or tokens are embedded into word vectors before entering the recurrent loops.The embeddings are represented as a sequence of vectors X = {x0, x1, . . . , xt, . . .}.
- Recurrent computation: The recurrent module uses update and reset gates to compute hidden-state representations.The prose identifies zt as the update gate, rt as the reset gate, ht as the hidden state, and h̃t as a new hidden state.
D. Differential Privacy
The method incorporates client-side randomized noise to protect data against inverse engineering while retaining compatibility with attentive aggregation.
- Privacy protection: Client-side randomized mechanisms provide differential privacy without revealing client data.The mechanism adds white noise with mean 0 and standard deviation σ to client parameters.
- Integration: FedAtt incorporates the differential-privacy mechanism before clients send updated parameters to the server.The implementation is written as server optimization for simplicity.
- Noise control: The coefficient β ∈ (0, 1] controls the effect of normal-noise randomization on federated aggregation performance.
IV. EXPERIMENTS
The experiments evaluate the proposed method on language-modeling datasets chosen to mimic decentralized mobile-keyboard applications.
- Experimental design: The study includes baseline comparisons and additional exploratory experiments assessing attentive aggregation performance.
A. Datasets
The study evaluates private neural language modeling on Penn Treebank, WikiText-2, and Reddit Comments, partitioning each dataset into client groups to mimic decentralized keyboard users.
- Datasets: Experiments use Penn Treebank, WikiText-2, and Reddit Comments to evaluate decentralized private neural language modeling for mobile keyboards.Reddit Comments are sampled from a large social-network dataset.
- Datasets: Reddit Comments are sampled at 1‰ to represent social-network data in the private language-modeling evaluation.
- Datasets: Training, validation, and testing token counts for the three datasets are reported in Table I.
- Data Partitioning: Each dataset is shuffled and randomly partitioned into 100 subsets representing 100 client users.Sampling is performed without replacement under an independently identical distribution.
B. Baselines and Settings
The experiments compare FedSGD, FedAvg, and the proposed FedAtt using testing perplexity after 50 communication rounds with small-scaled GRU client models.
- Baselines: The comparison includes FedSGD, FedAvg, and FedAtt, with FedAtt using attentive aggregation in a FedAvg-like setting.
- Evaluation Metric: Testing perplexity evaluates three federated methods across the three datasets and lower perplexity indicates better language-model prediction.
- Results: After 50 server-client communication rounds, FedAtt outperforms FedSGD and FedAvg in testing perplexity on all three datasets.For Penn Treebank and WikiText-2, the reported gains are significant at selected client fractions.
D. Client Fraction
Client participation affects both perplexity and communication efficiency. FedAtt remains below FedAvg across tested client fractions and can substantially reduce communication rounds under selected settings.
- Client Fraction: Testing perplexity fluctuates as the client fraction increases, and 70% client participation achieves the lowest perplexity in the reported experiment.
- Client Fraction: FedAtt achieves lower perplexity than FedAvg for every tested client-fraction setting.
- Communication Cost: Communication rounds are evaluated until testing perplexity falls below 90, using a small-scaled language model with 10% of clients selected.
- Communication Cost: FedAtt achieves comparable communication cost across the tested local-training epochs and batch sizes.
G. Scale of Model
The study compares model scales and tied versus untied embeddings under federated training. Tied embeddings generally improve testing perplexity, while attentive aggregation performs better than FedAvg in most scale settings.
- Cost considerations: Larger client models increase communication and computational costs when their parameter counts reach millions.The central server performs aggregation and optimization across communications with clients.
- Model scale: Embedding dimensions of 300, 650, and 1500 define the small, medium, and large client-model scales.Both tied and untied versions are evaluated; tied models require the RNN hidden state to match the embedding dimension.
- Model scale: Tied large-scale models achieve the best testing perplexity for both FedAvg and FedAtt.The comparison covers tied and untied models across small, medium, and large scales.
- Aggregation comparison: FedAtt achieves lower testing perplexity than FedAvg in four of six model-scale settings.The improvements occur for tied and untied small models, the tied medium model, and the tied large model.