Learning to Help in Multi-Class Settings

Authors: Yu Wu, Yansong Li, Zeyu Dong, Nitya Sathyavageeswaran, Anand D. Sarwate
Primary area: learning theory
Keywords: learning with abstention, learning with reject option, learning to defer, classification

Abstract:
Deploying complex machine learning models on resource-constrained devices is challenging due to limited computational power, memory, and model retrainability. To address these limitations, a hybrid system can be established by augmenting the local model with a server-side model, where samples are selectively deferred by a *rejector* and then sent to the server for processing. The hybrid system enables efficient use of computational resources while minimizing the overhead associated with server usage. The recently proposed Learning to Help (L2H) model proposed training a server model given a fixed local (client) model. This differs from the Learning to Defer (L2D) framework which trains the client for a fixed (expert) server. In both L2D and L2H, the training includes learning a rejector at the client to determine when to query the server. In this work, we extend the L2H model from binary to multi-class classification problems and demonstrate its applicability in a number of different scenarios of practical interest in which access to the server may be limited by cost, availability, or policy. We derive a stage-switching surrogate loss function that is differentiable, convex, and consistent with the Bayes rule corresponding to the 0-1 loss for the L2H model. Experiments show that our proposed methods offer an efficient and practical solution for multi-class classification in resource-constrained environments.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730473191842

thread_type:
official_review

<thread_object>
creator: reviewer_vq4j

<content>
summary:
This paper extends the L2H model from binary classification to multi-class settings, enabling collaboration between a local (client) model and a remote (server) model, particularly tailored for resource-constrained environments. A stage-switching surrogate loss function is developed to train the server model and a rejector, which facilitates selective sample deferral to the server, thereby optimizing overall performance.

strengths:
+ This paper contributes by extending the L2H framework to multi-class classification, which enhances the applicability of client-server cooperation in environments with limited resources.

+ The surrogate loss function supports asynchronous training, which aligns with scenarios of intermittent server access, beneficial for distributed learning settings.

+ This paper provides a solid theoretical foundation for the proposed method, supporting its validity and enhancing confidence in its application.

weaknesses:
- The client-server setup involves data offloading, raising potential privacy risks. The paper lacks a discussion on privacy-preserving methods, such as differential privacy, which are common in federated learning.

- While the surrogate loss aims to approximate the Bayes optimal classifier, the paper does not provide rigorous proofs for its consistency or convergence in multi-class settings, particularly under non-IID data distributions, a common issue in federated learning.

- The approach assumes IID data, which may not reflect real-world federated learning scenarios. Non-IID data could impair the rejector’s effectiveness in making deferral decisions, limiting the model's applicability.

- While the asynchronous training algorithm partially addresses intermittent server availability, it does not fully address the issues of model drift or maintain accuracy under unstable connectivity. Adaptive adjustments for synchronization frequency and performance are not fully explored.

- Relying on a fixed client classifier limits its ability to adapt to new data distributions or improve through collaborative training, potentially degrading performance over time and increasing dependency on the server model.

- The experimental setup involves relatively simple architectures (LeNet-5 for the client and AlexNet for the server). This limits the generalization of the framework to applications involving more complex models, such as those used in large-scale real-world systems. Additionally, the conclusion’s mention of "multi-party collaboration in hybrid machine learning systems, especially in the era of LLMs." could be misleading, as the paper does not provide any experiments or detailed discussion about large language models.

questions:
- Could the authors consider privacy-preserving techniques, such as differential privacy, to mitigate potential privacy risks? Additionally, how would these methods impact the proposed approach's effectiveness?

- Can the paper provide theoretical justifications for the surrogate loss's consistency and convergence properties, especially in non-IID data conditions?

- How could the framework adapt to handle non-IID data distributions across clients? Could the rejector's decision-making process be adjusted for such scenarios?

- For asynchronous training, could a method be introduced to adaptively adjust synchronization interval $S$, improving applicability to real-world deployments?

- Is there a way to make the client classifier adaptable to changing data distributions over time?

- Would the method be effective with more complex models, such as GPT-2, as hinted in the conclusion?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730533328306

thread_type:
official_review

<thread_object>
creator: reviewer_fjcz

<content>
summary:
This paper focuses on the collaboration between the client model and the server model by using a rejector. To train the rejector, client model, and server model, this paper derives a stage-switching surrogate loss function to allow the asynchronous training between the client model and the server model.

strengths:
1. The theoretical analysis of this paper is solid, which extends the binary classification setting of L2H to general multi-class settings.

2. The problem is important and the proposed method uses a surrogate loss to asynchronously train client and server model, which is simple yet theoretically effective.

weaknesses:
The experiments are limited. Experiments are mainly composed of hyper-parameter sensitive analysis evaluations, which lack comparison to baselines. For example, this paper should compare the proposed method with those using the synchronous training method or different surrogate loss functions in terms of overall performance.

questions:
1. What's the process during the inference stage with the surrogate loss?

2. What's the convergence rate difference between the asynchronous and synchronous methods?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731037907685

thread_type:
official_review

<thread_object>
creator: reviewer_aeuf

<content>
summary:
In this manuscript, a hybrid system is introduced, wherein the local model is supplemented by a server-side model and the rejector selects data samples to be directed to either the client model or the server model. The manuscript presents three distinct scenarios: “pay-per-request”, “intermittent availability”, and “bounded reject rate”, which align with practical constraints related to cost, availability, and policy, respectively. To address these challenges, a differentiable and convex stage-switching surrogate loss function is developed to iteratively optimize both the “rejector function” and “server model”. Experimental results validate the superior performance of the proposed framework.

strengths:
- The writing of this manuscript is clear, and easy to follow.
- This manuscript addresses three practical scenarios reflecting constraints related to cost, availability, and policy, and derives formal objective functions to represent them.

weaknesses:
- The authors state that they extend the “learning to help” framework to handle multi-class classification but do not explain the challenges involved in transitioning from binary-class to multi-class classification.
- In the Conclusion section, the authors state that the proposed framework opens new avenues for further exploration in multi-party collaboration. However, since the evaluation uses relatively simple networks, such as LeNet-5 and AlexNet, the authors should include more complex networks to thoroughly validate the framework's performance.
- Figure 2 illustrates that the hyper-parameters $ c_1$ and $c_2$ significantly affect the system's overall performance. Therefore, it is recommended that the authors evaluate the proposed framework's performance while varying $c_1$ and $c_2$ simultaneously. Additionally, the authors should provide guidance on achieving an optimal balance between $c_1$ and $c_2$.
- Table 1 shows that for the CIFAR-10 and SVHN datasets, when the rejector predicts “local”, the performance of $m(\cdot)$ is lower than $e(\cdot)$. Similarly, when the rejector predicts “remote”, $m(\cdot)$ still performs worse than $e(\cdot)$. Typically, the server model’s performance should surpass that of the local model when the rejector predicts “remote”. The authors should clarify this discrepancy.

questions:
As presented in the "Weaknesses" section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731168406417

thread_type:
official_review

<thread_object>
creator: reviewer_we9x

<content>
summary:
This paper extends the ``learning to help'' paradigm from binary classification to multi-class settings. Learning to help aims to learn a rejector for routing between edge legacy devices and server-side large models. This paper provides algorithms with theoretical analysis, followed by experiments on CIFAR-10 and SVHN.

strengths:
The setting of learning to help is quite interesting and meaningful. The authors extend it to multi-class classification cases, which have extensive applications. The proposed method is theoretically guaranteed. Also, the presentation is good.

weaknesses:
The main drawbacks of this paper may be in the experiment part.
- The datasets used are small, such as CIFAR-10 and SVHN. And there are only two datasets. Usually, evaluation on 3 datasets (or more) will be better.
- The used LeNet and AlexNet are tiny. Trending sota models like ViT can be considered.
- The learning to help scenarios can be very useful in large language model settings for on-device inference. Can the authors do some experiments on this? Or at least give a discussion.
- How do you partition the data for training the small legacy model and large server model? Is there any difference?
- I am just wondering about the costs of training the rejector and server model. Can the authors elaborate on this? If the costs are too much compared with the efficiency gains afterwards, why not use the server-side model solely?
- Can the authors do some case studies to show which samples need to be handled on the server instead of on the device? It will help to understand the system.
- Towards ``first principle'' to make the thing simple: for image classification tasks, one possible and simple solution is that during inference, check the prediction confidence (e.g., the softmax entropy) of the device model, and if the confidence is lower than a threshold (a predefined hyperparameter), send it to the server. Is this solution possible? It will not require any additional training for a rejector.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732208012506

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response: common points brought up by reviewers

comment:
We sincerely thank the reviewers for dedicating their time and providing such thoughtful and constructive feedback. Below we address some of the common themes and major points raised in the reviews. More specific responses will appear under the individual reviews. We believe that the points will provide some shared context for all reviewers.

### 1.  **L2H is not a form of federated learning (FL).** 

In particular, the target of L2H is to train an extra rejector and the server model to collaboratively work with the existing client model, while the target of FL is to train a single model in a distributed way. The three components local model, rejector, and the server model still engage in the test/deployment phase, while in FL only one model is needed to handle the tasks. More crucially, the dataset is not partitioned in L2H since the client model is fixed during the training phase, while each client model has its own training set on each client site in FL.

While it is also possible to do online training like the form of federated learning, eg. one server vs multiple clients with rejectors, *this would be a topic for future work*. This setting would raise many other practical issues common to all federated learning tasks: non-IID data, privacy, communication complexity, and so on.  As one example, it may be challenging to train a single server model to help multiple heterogeneous clients. As an analogy, it is also difficult for an instructor (=server) to teach (=help) students (=clients) with heterogeneous preparation/background knowledge.  

a. The L2H framework is designed to address scenarios in which the *client model cannot be updated once it is deployed*. This can happen if (for example) the client model is hard-coded into hardware. This assumption models many applications involving device-bound models in physical deployments such as “smart infrastructure” or automotive applications. 

Once deployed, it is very expensive to physically replace every device: a “safe streets” application may have cameras and computer vision systems at every intersection in a city: replacing all of the devices is like replacing every traffic light.   

One proposed architecture for these systems is to offload computation to a cloud service, but this can cause consistent latency for the task. 

The L2H model takes advantage of both solutions by letting a rejector only offload partial data to the server while keeping the use of the legacy model. We assume that the client models can be equipped with an updateable rejector so that the rejector and server models can be updated periodically. These updates will happen offline.

b. As with all ML systems, there is a training phase and an inference (deployment) phase for our system. In the training phase, the local model $m(x)$ is a fixed function whose parameters do not change. For each sample $x$, we calculate the loss $L\_1$ and update the parameters of the server model $e(x)$, and then calculate the loss $L\_2$ and update the parameters of reject function $r(x)$.  

c. For the inference phase, we do not use the loss function anymore. For each new sample $x$, the client will compute the trained $r(x)$ and decide whether to process $x$ locally with $m(x)$ or send it to the server who will process it with the trained $e(x)$.

d. The standard federated learning (FL) setting is to have multiple clients communicate with a server to train a single model. In this case, a client and server will hold the same model because the server pushes updated models to the clients. In L2H the client model is fixed and the server model is different. The assumption is that the server is more powerful than the client and can therefore have a better model. 

Ideally, we would be able to train a simple/lightweight rejector and server to help boost the performance of the overall system without requiring offloading of all test samples. 

(continued in next comment)
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper was reviewed by four experts in the field and received 6, 6, 8, 6 as the final ratings. The reviewers agreed that the setting of learning to help is interesting and meaningful, the proposed method is theoretically sound, the experimental results are encouraging, and that the paper is well-written and easy to follow. 

The reviewers raised concerns that the proposed method has been evaluated only on 2 datasets (CIFAR-10 and SVHN), with relatively simple architectures (LeNet and AlexNet). To address this, the authors have included experiments on the CIFAR-100 dataset with the ViT model in the revised version of the paper. The authors have also included a rigorous proof of the surrogate loss consistency in a multi-class setting. A table has also been included to show the impact of the hyper-parameters $c_{1}$ and $c_{2}$ on the system's overall performance using the CIFAR-100 dataset with the ViT model, in response to a reviewer's comment. Concerns were also raised about the cost of training, difference in convergence rate between the asynchronous and synchronous methods, and adaptability of the client classifier to changing data distributions over time, which were all addressed convincingly by the authors in the rebuttal.  

The reviewers, in general, have a positive opinion about the paper and its contributions (two of the reviewers had raised their scores from 5 to 6 after reviewing the author rebuttal). Based on the reviewers' feedback, the decision is to recommend the paper for acceptance to ICLR 2025. The reviewers have provided some valuable comments, such as the privacy risks in the client-server setup, comparison against additional baseline methods and evaluating the proposed method on larger datasets such as ImageNet. The authors are encouraged to address these in the final version of their paper. We congratulate the authors on the acceptance of their paper!