Importance-aware Co-teaching for Offline Model-based Optimization

Offline model-based optimization aims to find a design that maximizes a property of interest using only an offline dataset, with applications in robot, protein, and molecule design, among others. A prevalent approach is gradient ascent, where a proxy model is trained on the offline dataset and then used to optimize the design. This method suffers from an out-of-distribution issue, where the proxy is not accurate for unseen designs. To mitigate this issue, we explore using a pseudo-labeler to generate valuable data for fine-tuning the proxy. Specifically, we propose \textit{\textbf{I}mportance-aware \textbf{C}o-\textbf{T}eaching for Offline Model-based Optimization}~(\textbf{ICT}). This method maintains three symmetric proxies with their mean ensemble as the final proxy, and comprises two steps. The first step is \textit{pseudo-label-driven co-teaching}. In this step, one proxy is iteratively selected as the pseudo-labeler for designs near the current optimization point, generating pseudo-labeled data. Subsequently, a co-teaching process identifies small-loss samples as valuable data and exchanges them between the other two proxies for fine-tuning, promoting knowledge transfer. This procedure is repeated three times, with a different proxy chosen as the pseudo-labeler each time, ultimately enhancing the ensemble performance. To further improve accuracy of pseudo-labels, we perform a secondary step of \textit{meta-learning-based sample reweighting}, which assigns importance weights to samples in the pseudo-labeled dataset and updates them via meta-learning. ICT achieves state-of-the-art results across multiple design-bench tasks, achieving the best mean rank of $3.1$ and median rank of $2$, among $15$ methods. Our source code can be found here.

Paper

Similar papers

Peer review

Reviewer DRKm7/10 · confidence 4/52023-07-07

Summary

This paper proposes a novel model-based optimization method that combines ensemble-based ideas and meta learning. The authors introduce three proxy models, each capable of generating pseudo-labels for the training of the other models, effectively making them co-teachers. This co-teaching process serves as an ensemble training method, mitigating potential risks associated with pseudo-labeling. Additionally, the paper incorporates meta learning techniques by re-weighting the samples, assigning higher weights to accurate samples. The intuitive nature of this idea and the reasonable performance observed in the experiments make a compelling case for accepting this paper.

Strengths

1. There are several methods available for stabilizing the learning process with pseudo-labels, and this idea has the potential to be applied to various literature in the field. 2. The idea presented in this paper is straightforward and effectively communicated, allowing readers to grasp its essence easily. 3. One of the key contributions of this paper is the thorough analysis of hyperparameters, which play a critical role in model-based optimization. The authors provide valuable insights into optimizing these parameters to achieve better performance.

Weaknesses

1. Although this paper presents a highly intuitive method, it lacks a detailed mathematical explanation for why this approach is effective. It would be beneficial to include analysis based on experiments to provide more insight into the working mechanism. 2. It appears that several ideas in this paper draw inspiration from the prior method of BDI [1], particularly in the context of offline model-based optimization. While this viewpoint is intriguing, a more explicit analysis comparing the proposed method with BDI would greatly assist readers in understanding the similarities and differences. 3. It would be valuable to reference offline biological design methods [2,3], even though they may not directly address continuous tasks. Specifically, [2] can be good reference to discuss this method due to their utilization of pseudo-labeling and their success in stabilizing the bootstrapping process, which aligns with the idea presented in this paper. Including a dedicated discussion section on these references would greatly benefit future researchers. [1] Chen, Can, et al. "Bidirectional learning for offline infinite-width model-based optimization." Advances in Neural Information Processing Systems 35 (2022): 29454-29467. [2] Kim, Minsu, et al. "Bootstrapped Training of Score-Conditioned Generator for Offline Design of Biological Sequences." arXiv preprint arXiv:2306.03111 (2023). [3] Jain, Moksh, et al. "Biological sequence design with gflownets." International Conference on Machine Learning. PMLR, 2022.

Questions

1. The paper employs three proxy co-teachers, but it would be interesting to investigate the impact of reducing the number to just two proxies. Are there any notable observations or differences in performance with this modification? 2. The meta learning process appears to involve bi-level optimization of parameters and weights. Considering the potential complexity of this approach, it raises concerns about computational resources required to execute the code. It would be helpful to provide insights or strategies to address this issue. 3. While this method demonstrates impressive performance in continuous design tasks, its effectiveness in discrete tasks, such as TFBind8, seems to be limited. How does this method perform on other protein tasks like GFP or UTR? Is there any intuition as to why this method excels specifically in continuous tasks? 4. The method's performance, when measured by the median score (50th percentile), appears to be weaker in comparison to the maximum score (100th percentile). Can you provide some intuition or explanation for this observation? It would be valuable to understand the factors contributing to this difference in performance.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

This paper really already gives limitations which I really appreciate it.

Reviewer VPeC6/10 · confidence 3/52023-07-08

Summary

The paper proposes a new method to address the out-of-distribution issue of offline model-based optimization. The author first generates a pseudo-labeled dataset and identifies valuable (important) data for proxy co-teaching. To alleviate potential inaccuracies in pseudo-labels, a meta-learning framework is proposed to adjust the importance weights of samples. Experiments on both continuous and discrete tasks demonstrate the effectiveness of the proposed method.

Strengths

1. The out-of-distribution issue in offline model-based optimization is a challenging problem. The proposed method is well-motivated, technically sound, and interesting. 2. Experiments and ablation studies on both continuous and discrete tasks demonstrate the effectiveness of the proposed method. 3. The paper is well-written and easy to follow.

Weaknesses

1. In Lines 150-151, the authors claim that "small-loss samples typically contain valuable knowledge, making them ideal for enhancing proxy robustness". However, in my opinion, these small-loss samples might be some samples easy to learn since different proxies generate similar outputs for them. These easy samples may not be the truly valuable samples that benefit model generalization. 2. In Section 3.2, a meta-learning framework is proposed to adjust the importance weights of samples. However, there still exist distributional shifts between the inner-loop and outer-loop. Therefore, the loss in Equation 7 may not be informative enough to identify truly important samples for generalization. Samples with hight weight might be the ones similar to the offline dataset. Therefore, the proposed method may not be applicable when the distributional shift is large.

Questions

Is it possible to visualize the so-called valuable samples such that the intuitive interpretation can be obtained?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes.

Reviewer P8Ru5/10 · confidence 3/52023-07-11

Summary

The paper addresses the offline MBO problem, which aims to synthesize an instance $x$ that maximizes the proxy $f(\cdot)$ estimated by a neural network. The popular approach, gradient ascent, updates the instance $x_{t+1}$ to have a higher proxy value compared to the previous instance $x_t$. However, this approach encounters a challenge where the proxy $f(\cdot)$ becomes less accurate as the iterations progress, as the instance $x_t$ becomes more out-of-distribution. To overcome this challenge, the paper proposes two techniques inspired by the field of noisy labels. Firstly, the paper applies co-teaching, a method that filters confident samples by selecting those consistent across two models. In this case, since the labels are generated by a neural network, the final method involves three models: one pseudo-labeler and two co-teachers. The three models are trained alternately, changing their roles. Secondly, to further enhance performance, the paper introduces sample reweighing, a technique that assigns higher weights to confident samples using meta-learning. When combined, the proposed ICT (importance-aware co-teaching) method outperforms previous approaches.

Strengths

- The idea of applying noisy label techniques to handle pseudo-labels of OOD samples makes sense. - The proposed two components are necessary, as shown in the ablation study.

Weaknesses

### Experiments are not convincing My biggest concern is the experimental results. The gains over the prior works are marginal compared to the high variances. The rank seems better at first glance, but it is not really convincing as the gap of proxy scores is too incremental. ### On the evaluation metric I know this is not the problem of this paper but of the entire field. Still, reporting only the proxy value as the evaluation metric may not be enough. How can one confirm that the updated instance $x_T$ follows the valid form of $\mathcal{X}$? ### Limited technical novelty The proposed method is basically an application of known methods such as co-teaching and sample reweighing. Although the three model parts were somewhat new, thanks to the use of the pseudo-labeler, the technical innovation is not sufficiently strong. ### Presentation can be polished The writing contains various repetitions, with similar words appearing in the abstract, introduction, method, and conclusion. It would be great if the paper could provide more information and prune redundant details. Additionally, Figure 1 and 2 are hard to understand without reading the full texts. The captions can be polished to be self-contained.

Questions

Written in the weakness.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

Discussed.

Reviewer j29o6/10 · confidence 4/52023-07-23

Summary

The work proposes an Importance-aware Co-Teaching for Offline Model-based Optimization (ICT) as a solution for Offline model-based optimization problems. The ICT algorithm is comprised of 2 steps: (i) pseudo-label-driven co-teaching and (ii) per-sample reweighting. Overall the approach aims to learn an ensemble of methods where the targets predicted by one is treated as the pseudo-label for the other models. Furthermore, the more confident samples with low losses are used to give feedback to the proxy model.

Strengths

* The paper proposes their intuition backed by two clearly laid out steps to walk the readers through the working of the model. * The extensive results on multiple datasets of both continuous and discrete domains help better bolster their approach. * The ablation of each step is useful to better weigh each design choice of ICT

Weaknesses

* Some of the results in Table 1, Table 2 are within standard deviation. I would advice the authors to also perform a p-test to better strengthen their claims. * It would be nice to expand more on the per-sample weighting part. My knowledge on this subject often uses a validation data to as a proxy for minimizing loss, similar to Meta-weight-Net [1]. In the current training regime, it is not clear to me why the model cannot converge to the trivial solution of a zero vector as per-sample weights. [1] Shu, Jun, et al. "Meta-weight-net: Learning an explicit mapping for sample weighting." Advances in neural information processing systems 32 (2019).

Questions

I am interested in understand two design choices taken in this work: * In co-teaching, the two models other than proxy share knowledge about top-K points which both have learnt very well. If the K points are in-distribution and all models have learnt them perfectly, I see little to gain from this step. I was thinking using the worst-K points might be more interesting, as I am not sure how the training dynamics would play out there. * I per-sample weighting, i am curious how the current proposed regime fares compared to other solutions such as RGD [2] which have been shown to achieve improved generalization inspired by DRO literature. RGD would also remove the need to learn another meta-network However, this again focusses on giving more importance to samples the model is weak at, and roughly connects back to the first point. [2] Kumar, Ramnath, et al. "Stochastic Re-weighted Gradient Descent via Distributionally Robust Optimization." arXiv preprint arXiv:2306.09222 (2023).

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The approach is in the domain of optimization, and does not directly have any negative impact by itself. The authors have provided a nice paragraph covering this aspect in their paper.

Reviewer j29o2023-08-11

Thank you for your detailed response. I have read through most of the comments by other reviewers, as well as the rebuttal. This clears most of my questions and concerns. Thus, I have increased my score.

Authorsrebuttal2023-08-11

Thank you very much for your prompt feedback

Thank you for your prompt feedback. We appreciate your review of our responses and others' comments. It's encouraging to see that our clarifications resolved your concerns, and we will make the necessary revisions to the paper as discussed.

Reviewer P8Ru2023-08-13

Concerns addressed

Thank you for the rebuttal. After reading other reviews and responses, I have increased my rating to acceptance. It would be greatly appreciated if the authors could include the responses, such as clarifications on the MBO-specific novelties over the prior co-teaching, in the revised manuscript. This could be accommodated in the additional space created by reducing redundancies.

Authorsrebuttal2023-08-13

Thanks for your feedback

Thank you for your feedback and the adjusted score. We'll ensure that clarifications on MBO-specific novelties, along with other responses, are included in the final manuscript, utilizing the space created by reducing redundancies.

Reviewer DRKm2023-08-17

Thank you for the rebuttal. I keep my score supporting this paper to be accepted.

Authorsrebuttal2023-08-17

Thank you

Thank you for your support and for maintaining your score for accepting our paper. We greatly appreciate your valuable feedback, and we will diligently revise the paper as we discuss.

Authorsrebuttal2023-08-18

Look forward to hearing from you

Thank you for your comprehensive review and insightful feedback. We've addressed your concerns as follows: 1. **Small-Loss Samples**: Clarified the value of small-loss samples in enhancing our model's robustness in noisy label contexts. 2. **Meta-Learning Framework**: Acknowledged the persistent issue of large distributional shifts in offline MBO, yet demonstrated our method's effectiveness with extensive experiments. 3. **Visualization of Valuable Samples**: Adopted an analytical approach to identify valuable samples, emphasizing pseudo-label accuracy. Could you please clarify if there are any unresolved issues? We have engaged constructively with the other three reviewers over the past week and eagerly anticipate your further insights to continue refining our work. Thank you, and we look forward to hearing from you.

Reviewer VPeC2023-08-18

Thank the authors for addressing my concerns. I intend to agree that the small-loss samples might be the clean ones in noisy label scenarios, and therefore they are beneficial for model robustness. I also checked out the comments of other reviewers and I do not have any major concerns. I prefer to increase my score and vote for acceptance.

Authorsrebuttal2023-08-18

Thanks

Thank you for your thoughtful review and your consideration of the adjustments we have made in response to your insights. We're delighted to hear that you recognize the merits of our approach, particularly with regards to the use of small-loss samples in noisy label scenarios. We sincerely appreciate your increased score 6 and vote for acceptance. Your feedback has been instrumental in enhancing our work, and we are grateful for your positive assessment.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC