SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Assumption 4.3. For all  $(i,j)\in S_1\times S_2$ , we assume that
Algorithm 2 Stochastic Optimization algorithm for Non-smooth TCCO (SONT)
1: Initialization:  $\mathbf{w}_0$ ,  $\{u_{i,0}:i\in S_1\}$ ,  $v_{i,j,0} = h_{i,j}(\mathbf{w}_0;\mathcal{B}_{3,i,j}^0)$  for all  $(i,j)\in S_1\times S_2$
2: for  $t = 0,\dots ,T - 1$  do
3: Sample batches  $\mathcal{B}_1^t\subset \mathcal{S}_1$ $\mathcal{B}_2^t\subset \mathcal{S}_2$  , and  $\mathcal{B}_{3,i,j}^{t}\subset \mathcal{D}_{i,j}$  for  $i\in \mathcal{B}_1^t$  and  $j\in \mathcal{B}_2^t$
$$
4: \quad v _ {i, j, t + 1} = \left\{ \begin{array}{l l} \Pi_ {\tilde {C} _ {h}} [ (1 - \tau_ {1}) v _ {i, j, t} + \tau_ {1} h _ {i, j} (\mathbf {w} _ {t}; \mathcal {B} _ {3, i, j} ^ {t}) + \gamma_ {1} (h _ {i, j} (\mathbf {w} _ {t}; \mathcal {B} _ {3, i, j} ^ {t}) - h _ {i, j} (\mathbf {w} _ {t - 1}; \mathcal {B} _ {3, i, j} ^ {t})) ], \\ v _ {i, j, t}, \quad (i, j) \notin \mathcal {B} _ {1} ^ {t} \times \mathcal {B} _ {2} ^ {t} \end{array} \right.
$$
$$
\begin{array}{l} 5: \quad u _ {i, t + 1} = \left\{ \begin{array}{l} (1 - \tau_ {2}) u _ {i, t} + \frac {1}{B _ {2}} \sum_ {j \in \mathcal {B} _ {2} ^ {t}} [ \tau_ {2} g _ {i} (v _ {i, j, t}) + \gamma_ {2} (g _ {i} (v _ {i, j, t}) - g _ {i} (v _ {i, j, t - 1}) ], \quad i \in \mathcal {B} _ {1} ^ {t} \\ u _ {i, t}, \quad i \notin \mathcal {B} _ {1} ^ {t} \end{array} \right. \\ 6: \quad G _ {t} = \frac {1}{B _ {1}} \sum_ {i \in \mathcal {B} _ {1} ^ {t}} \left[ \left(\frac {1}{B _ {2}} \sum_ {i \in \mathcal {B} _ {2} ^ {t}} \nabla h _ {i, j} (\mathbf {w} _ {t}; \mathcal {B} _ {3, i, j} ^ {t}) \partial g _ {i} (v _ {i, j, t})\right) \partial f _ {i} (u _ {i, t}) \right] \\ 7: \quad \text {U p d a t e} \mathbf {w} _ {t + 1} = \mathbf {w} _ {t} - \eta G _ {t} \\ \end{array}
$$
-  $f_{i}$  is  $C_f$ -Lipschitz continuous,  $\rho_f$ -weakly-convex and non-decreasing;
-  $g_{i}$  is  $\rho_{g}$ -weakly-convex and  $C_{g}$ -Lipschitz continuous.  $h_{i,j}(\cdot ;\xi)$  is  $C_{h}$ -Lipschitz continuous.
- Either  $g_{i}$  is non-decreasing,  $h_{i,j}$  is  $L_{h}$ -weakly-convex or  $g_{i}$  is monotone,  $h_{i,j}$  is  $L_{h}$ -smooth.
- Stochastic estimators  $h_{i,j}(\mathbf{w},\xi)$  and  $\partial h_{i,j}(\mathbf{w},\xi)$  have bounded variance  $\sigma^2$ , and  $\| h_{i,j}(\mathbf{w})\| \leq \tilde{C}_h$ .  $\mathbb{E}_i\| g_i(v) - \frac{1}{n_2}\sum_{j\in S_2}g_i(v)\| ^2\leq \sigma^2$  for any  $v$ .
The weak convexity of  $F(\mathbf{w})$  in (2) is guaranteed by the following Proposition.
Proposition 4.4. Under Assumption 4.3,  $F(\mathbf{w})$  in (2) is  $\rho_{F}$ -weakly-convex with  $\rho_{F} = \sqrt{d_{1}} (\sqrt{d_{2}} L_{h}C_{g} + \rho_{g}C_{h}^{2})C_{f} + \rho_{f}C_{g}^{2}C_{h}^{2}$ .
We extend SONX to Algorithm 2 for (2), which is referred to as SONT. For dealing with the extra layer of compositional problem, we maintain another multi-block variable to track the extra layer of function value estimation. To understand this, we first write down the true subgradient:
$$
\partial F (\mathbf {w}) = \frac {1}{n _ {1}} \sum_ {i \in \mathcal {S} _ {1}} \left[ \left(\frac {1}{n _ {2}} \sum_ {j \in \mathcal {S} _ {2}} \nabla h _ {i, j} (\mathbf {w}) \partial g _ {i} (h _ {i, j} (\mathbf {w}))\right) \partial f _ {i} \left(\frac {1}{n _ {2}} \sum_ {j \in \mathcal {S} _ {2}} g _ {i} (h _ {i, j} (\mathbf {w}))\right) \right].
$$
To approximate this subgradient, we need the estimations of  $\frac{1}{n_2}\sum_{j\in S_2}g_i(h_{i,j}(\mathbf{w}))$  and  $h_{i,j}(\mathbf{w})$  which can be tracked by using MSVR estimators denoted by  $\{u_{i,t}:i\in S_1\}$  and  $\{v_{i,j,t}:(i,j)\in S_1\times S_2\}$ , respectively. As a result, a stochastic estimation of  $\partial F(\mathbf{w}_t)$  is computed in step 6 of Algorithm 2, and the model parameter is updated similarly as before.
In this section, we present the proof sketch of the convergence guarantee for Algorithm 1. The analysis for Algorithm 2 follows in a similar manner. The detailed proofs can be found in Appendix A (please refer to the supplement). Before starting the proof, we define a constant  $M^2 \geq C_f^2 C_g^2$  so that under Assumption 4.1 we have  $\mathbb{E}_t[\| G_t\|^2] \leq M^2$ . Then we start by giving the error bound of the MSVR estimator in Algorithm 1. The following norm bound of the estimation error follows from the squared-norm error bound in Lemma 1 from [15], whose proof is given in Appendix D.3.
Lemma 4.5. Consider the update for  $\{u_{i,t}:i\in S\}$  in Algorithm 1. Assume  $g_{i}$  is  $C_g$ -Lipschitz for all  $i\in S$ . With  $\gamma = \frac{n - B_1}{B_1(1 - \tau)} + (1 - \tau)$ ,  $\tau \leq \frac{1}{2}$ , we have
$$
\mathbb {E} \left[ \frac {1}{n} \sum_ {i \in \mathcal {S}} \| u _ {i, t + 1} - g _ {i} (\mathbf {w} _ {t + 1}) \| \right] \leq (1 - \frac {B _ {1} \tau}{2 n}) ^ {t + 1} \frac {1}{n} \sum_ {i \in \mathcal {S}} \| u _ {i, 0} - g _ {i} (\mathbf {w} _ {0}) \| + \frac {2 \tau^ {1 / 2} \sigma}{B _ {2} ^ {1 / 2}} + \frac {4 n C _ {g} M \eta}{B _ {1} \tau^ {1 / 2}}.
$$
For simplicity, denote by  $\hat{\mathbf{w}}_t\coloneqq \mathrm{prox}_{F / \bar{\rho}}(\mathbf{w}_t)$ . Then using the definition of Moreau envelope and the update rule of  $\mathbf{w}_t$ , we can obtain a bound for the change in the Moreau envelope,
$$
\mathbb {E} _ {t} \left[ F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t + 1}\right) \right] \leq F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t}\right) + \bar {\rho} \eta \langle \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}, \mathbb {E} _ {t} \left[ G _ {t} \right] \rangle + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2}. \tag {5}
$$
where  $\mathbb{E}_t[G_t] = \frac{1}{n}\sum_{i\in S_1}\partial g_i(\mathbf{w}_t)\partial f_i(u_{i,t})$  is the subgradient approximation based on the MSVR estimator  $u_{i,t}$  of the inner function value. This is a standard result in weakly-convex optimization [6].
To bound the inner product  $\langle \hat{\mathbf{w}}_t - \mathbf{w}_t, \mathbb{E}_t[G_t] \rangle$  on the right-hand-side of (5), we apply the assumptions that  $f_i$  is weakly-convex, Lipschitz continuous and non-decreasing, and  $g_i$  is weakly-convex. Its upper bound is given as follows.
$$
\begin{array}{l} \left(\hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t}\right) ^ {\top} \mathbb {E} _ {t} \left[ G _ {t} \right] \leq F (\hat {\mathbf {w}} _ {t}) - F (\mathbf {w} _ {t}) + \frac {1}{n} \sum_ {i \in \mathcal {S}} \left[ f _ {i} \left(g _ {i} \left(\mathbf {w} _ {t}\right)\right) - f \left(u _ {i, t}\right) - \partial f \left(u _ {i, t}\right) ^ {\top} \left(g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t}\right) \right. \\ + \rho_ {f} \| g _ {i} (\mathbf {w} _ {t}) - u _ {i, t} \| ^ {2} + \left(\frac {\rho_ {g} C _ {f}}{2} + \rho_ {f} C _ {g} ^ {2}\right) \| \hat {\mathbf {w}} _ {t} - \mathbf {w} _ {t} \| ^ {2} ]. \tag {6} \\ \end{array}
$$
Due to the  $\rho_F$ -weak convexity of  $F(\mathbf{w})$ , we have  $(\bar{\rho} - \rho_F)$ -strong convexity of  $\mathbf{w} \mapsto F(\mathbf{w}) + \frac{\bar{\rho}}{2} \| \mathbf{w}_t - \mathbf{w} \|^2$ . Then it follows  $F(\hat{\mathbf{w}}_t) - F(\mathbf{w}_t) \leq \left( \frac{\rho_F}{2} - \bar{\rho} \right) \| \mathbf{w}_t - \hat{\mathbf{w}}_t \|^2$ . Combining this with inequalities (5), (6), and setting  $\bar{\rho}$  sufficiently large we have
$$
\begin{array}{l} \mathbb {E} _ {t} \left[ F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t + 1}\right) \right] \leq F _ {1 / \bar {\rho}} \left(\mathbf {w} _ {t}\right) + \frac {\eta^ {2} \bar {\rho} M ^ {2}}{2} + \frac {\bar {\rho} \eta}{n} \sum_ {i \in S} \left[ - \frac {\bar {\rho}}{2} \| \mathbf {w} _ {t} - \hat {\mathbf {w}} _ {t} \| ^ {2} \right. \tag {7} \\ + f _ {i} \left(g _ {i} \left(\mathbf {w} _ {t}\right)\right) - f \left(u _ {i, t}\right) - \partial f _ {i} \left(u _ {i, t}\right) ^ {\top} \left(g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t}\right) + \rho_ {f} \| g _ {i} \left(\mathbf {w} _ {t}\right) - u _ {i, t} \| ^ {2} ]. \\ \end{array}
$$
Recall Lemma 3.2, we have  $\| \mathbf{w}_t - \hat{\mathbf{w}}_t\|^2 = \frac{1}{\bar{\rho}^2}\|\nabla F_{1/\bar{\rho}}(\mathbf{w}_t)\|^2$ . Moreover, the last three terms on the R.H.S of inequality (7) can be bounded using the Lipschitz continuity of  $f_i$  and the error bound given in Lemma 4.5. Then we can conclude the complexity of SONX with the following theorem.
Theorem 4.6. Under Assumption 4.1 with  $\gamma = \frac{n - B_1}{B_1(1 - \tau)} + (1 - \tau)$ ,  $\tau = \mathcal{O}(B_2\epsilon^4) \leq \frac{1}{2}$ ,  $\eta = \mathcal{O}\left(\frac{B_1B_2^{1/2}\epsilon^4}{n}\right)$ , and  $\bar{\rho} = \rho_F + \rho_gC_f + 2\rho_fC_g^2$ , Algorithm 1 converges to an  $\epsilon$ -stationary point of the Moreau envelope  $F_{1/\bar{\rho}}$  in  $T = \mathcal{O}\left(\frac{n}{B_1B_2^{1/2}}\epsilon^{-6}\right)$  iterations.
Remark. Similar to the complexity for smooth FCCO problems [26, 15], Theorem 4.6 guarantees that SONX for NSWC FCCO has a parallel speed-up in terms of the batch size  $B_{1}$  and linear dependency on  $n$ . The dependency of the complexity on the batch size  $B_{2}$  is due to the use of MSVR estimator, which matches the results in [15]. If the MSVR estimator in SONX is replaced by moving average estimator, the complexity becomes  $\mathcal{O}\left(\frac{n}{B_1B_2}\epsilon^{-8}\right)$  (cf. Appendix B).
Following a similar proof strategy, the convergence guarantee of Algorithm 2 is given below.
Theorem 4.7. (Informal) Under Assumption 4.3, with appropriate values of  $\gamma_1, \gamma_2, \tau_1, \tau_2, \eta$  and a proper constant  $\bar{\rho}$ , Algorithm 2 converges to an  $\epsilon$ -stationary point of the Moreau envelope  $F_{1 / \bar{\rho}}$  in  $T = \mathcal{O}\left(\max \left\{\frac{1}{B_3^{1/2}}, \frac{n_1^{1/4}}{B_1^{1/4}n_2^{1/4}}, \frac{n_1^{1/2}}{B_1^{1/2}n_2^{1/2}}\right\} \frac{n_1n_2}{B_1B_2}\epsilon^{-6}\right)$  iterations.
Remark. In the worst case, the complexity has a worse dependency on  $n_1 / B_1$ , i.e.,  $\mathcal{O}(n_1^{3/2} / B_1^{3/2})$ . This is caused by the two layers of block-sampling update for  $\{u_{i,t}, i \in S_1\}$  and  $\{v_{i,j,t} : (i,j) \in S_1 \times S_2\}$ . When  $n_1 = B_1 = 1$  and  $B_3 \leq \sqrt{n_2}$ , the complexity of SONT becomes similar as SONX, which is understandable as the inner two levels in TCCO are the same as FCCO.
NSWC FCCO finds important applications in group distributionally robust optimization (group DRO) and two-way partial AUC (TPAUC) maximization.
Consider  $N$  groups with different distributions. Each group  $k$  has an averaged loss  $L_{k}(w) = \frac{1}{n_{k}}\sum_{i=1}^{n_{k}}\ell(f_{w}(x_{i}^{k}),y_{i}^{k})$ , where  $w$  is the model parameter and  $(x_{i}^{k},y_{i}^{k})$  is a data point. It has been shown in previous study [23] that the group DRO problem can be formulated into
$$
\min  _ {w} \min  _ {s} F (w, s) = \frac {1}{K} \sum_ {k = 1} ^ {N} [ L _ {k} (w) - s ] _ {+} + s.
$$
This formulation can be mapped into non-smooth weakly-convex FCCO under certain assumptions. Due to space limitation, we defer the comprehensive discussion of group DRO to Appendix E. The rest of this section focuses on TPAUC maximization.
Let  $X$  denote an input example and  $h_{\mathbf{w}}(X)$  denote a prediction of a parameterized deep net on data  $X$ . Denote by  $S_{+}$  the set of  $n_{+}$  positive examples and by  $S_{-}$  the set of  $n_{-}$  negative examples. TPAUC measures the area under ROC curve where the true positive rate (TPR) is higher than  $\alpha$  and the false positive rate (FPR) is lower than an upper bound  $\beta$ . A surrogate loss for optimizing TPAUC
with  $\mathrm{FPR} \leq \beta$ ,  $\mathrm{TPR} \geq \alpha$  is given by [34]:
$$
\min  _ {\mathbf {w}} \frac {1}{n _ {+}} \frac {1}{n _ {-}} \sum_ {X _ {i} \in \mathcal {S} _ {+} ^ {\uparrow} [ 1, k _ {1} ]} \sum_ {X _ {j} \in \mathcal {S} _ {-} ^ {\downarrow} [ 1, k _ {2} ]} \ell \left(h _ {\mathbf {w}} \left(X _ {j}\right) - h _ {\mathbf {w}} \left(X _ {i}\right)\right), \tag {8}
$$
where  $\ell(\cdot)$  is a convex, monotonically non-decreasing surrogate loss of the indicator function  $\mathbb{I}(h_{\mathbf{w}}(X_j) \geq h_{\mathbf{w}}(X_i))$ ,  $S_+^\uparrow [1, k_1]$  is the set of positive examples with  $k_1 = \lfloor n_+ \alpha \rfloor$  smallest scores, and  $S_-^\downarrow [1, k_2]$  is the set of negative examples with  $k_2 = \lfloor n_- \beta \rfloor$  largest scores. To tackle the challenge of selecting examples from  $S_+^\uparrow [1, k_1]$  and  $S_-^\downarrow [1, k_2]$ , the above problem is cast into the following [44]:
$$
\min  _ {\mathbf {w}, \mathbf {s} ^ {\prime}, \mathbf {s}} \frac {1}{n _ {+}} \sum_ {X _ {i} \in \mathcal {S} _ {+}} f _ {i} \left(\psi_ {i} \left(\mathbf {w}, s _ {i}\right), s ^ {\prime}\right), \tag {9}
$$
where  $f_{i}(g,s^{\prime}) = s^{\prime} + \frac{(g - s^{\prime})_{+}}{\alpha}$ ,  $\psi_{i}(\mathbf{w},s_{i}) = \frac{1}{n_{-}}\sum_{X_{j}\in \mathcal{S}_{-}}s_{i} + \frac{(\ell(h_{\mathbf{w}}(X_{j}) - h_{\mathbf{w}}(X_{i})) - s_{i})_{+}}{\beta}$ ,
where  $\mathbf{s} = (s_1, \ldots, s_{n+})$ . We will consider two scenarios, namely regular learning scenario where  $X_i \in \mathbb{R}^{d_0}$  is an instance, and multi-instance learning (MIL) scenario where  $X_i = \{\mathbf{x}_i^1, \ldots, \mathbf{x}_i^{m_i} \in \mathbb{R}^{d_0}\}$  contains multiple instances (e.g., one patient has hundreds of high-resolution CT images). A challenge in MIL is that the number of instances  $m_i$  for each data might be large such that it is difficult to load all instances into the memory for mini-batch training. It becomes more nuanced especially because MIL involves a pooling operation that aggregates the predicted information of individual instances into a single prediction, which can be usually written as a compositional function with the inner function being an average over instances from  $X$ . For simplicity of exposition, below we consider the mean pooling  $h_{\mathbf{w}}(X) = \frac{1}{|X|} \sum_{\mathbf{x} \in X} e(\mathbf{w}_e; \mathbf{x})^\top \mathbf{w}_c$ , where  $e(\mathbf{w}_e, \mathbf{x})$  is the encoded feature representation of instance  $\mathbf{x}$  with a parameter  $\mathbf{w}_e$ , and  $\mathbf{w}_c$  is the parameter of the classifier. We will map the regular learning problem as NSWC FCCO and the MIL problem as NSWC TCCO.
The problem (9) is slightly more complicated than (1) or (2) due to the presence of  $s'$ ,  $\mathbf{s}$ . In order to understand the applicability of our analysis and results to (9), we ignore  $s'$ ,  $\mathbf{s}$  for a moment. In the regular learning setting when  $h_{\mathbf{w}}(X) = e(\mathbf{w}_e, X)^\top \mathbf{w}_c$  can be directly computed, we can map the problem into NSWC FCCO, where  $f_i(g, s')$  is non-smooth, convex, and non-decreasing in terms of  $g$ , and  $g_i(\mathbf{w}, s_i) = \psi_i(\mathbf{w}, s_i)$  is non-smooth, and is proved to be weakly when  $\ell(\cdot)$  is convex and  $h_{\mathbf{w}}(X)$  is smooth in terms of  $\mathbf{w}$ . In the MIL setting with mean pooling, we can map the problem into NSWC TCCO by defining  $h_i(\mathbf{w}) = \frac{1}{|X_i|} \sum_{\mathbf{x} \in X_i} e(\mathbf{w}_e; \mathbf{x})^\top \mathbf{w}_c$ ,  $h_{ij}(\mathbf{w}) = h_j(\mathbf{w}) - h_i(\mathbf{w})$  and  $g_i(h_{i,j}(\mathbf{w}), s_i) = s_i + \frac{(\ell(h_{i,j}(\mathbf{w})) - s_i)_+}{\beta}$ , and  $f_i(g_i, s') = s' + \frac{(g_i - s')_+}{\alpha}$ , where  $f_i$  is non-smooth, convex, and non-decreasing in terms of  $g_i$ , and  $g_i(h_{ij}(\mathbf{w}), s_i)$  is non-smooth, convex, monotonic in terms of  $h_{ij}(\mathbf{w})$  when  $\ell(\cdot)$  is convex and monotonically non-decreasing, and  $g_i(h_{ij}(\mathbf{w}), s_i)$  is weakly convex in terms of  $\mathbf{w}$  when  $h_{ij}(\mathbf{w})$  is smooth and Lipschitz continuous in terms of  $\mathbf{w}$ . Hence, the problem (9) satisfies the conditions in Assumption 4.1 for the regular learning setting and that in Assumption 4.3 for the MIL with mean pooling under mild regularity conditions of the neural network. We present full details in Appendix C.1 for interested readers.
To compute the gradient estimator w.r.t  $\mathbf{w}$ ,  $u_{i,t}$  will be maintained for tracking  $g_{i}(\mathbf{w}, s_{i})$  in the regular setting or  $\frac{1}{n-}\sum_{X_{j} \in S_{-}} g_{i}(h_{i,j}(\mathbf{w}), s_{i})$  in the MIL setting,  $v_{i,t}$  will be maintained for tracking  $h_{i}(\mathbf{w})$  in the MIL setting, which are updated similar to that in SONX and SONT. One difference from SONT is that  $v_{i,j,t}$  is decoupled into  $v_{i,t}$  and  $v_{j,t}$  due to that  $h_{i,j}$  can be decoupled. In terms of the extra variable  $s'$ ,  $\mathbf{s}$ , the objective function is convex w.r.t both  $s'$  and  $\mathbf{s}$ , which allows us to simply update  $s'$  by SGD using the stochastic gradient estimator  $\frac{1}{B_1} \sum_{i \in \mathcal{B}_1^t} \partial_{s'} f_i(u_{i,t}, s_t')$  and we update  $s_i$  by SGD using the stochastic gradient estimator  $\left[\frac{1}{B_2} \sum_{j \in \mathcal{B}_2^t} \partial_{s_i} g_i(v_{j,t} - v_{i,t}, s_{i,t})\right] \partial_u f_i(u_{i,t}, s_t')$ . Detailed updates are presented in Algorithm 5 and Algorithm 6 in Appendix C.2. We can extend the convergence analysis of SONX and SONT to the two learning settings of TPAUC maximization, which is included in Appendix C.4. Finally, it is worth mentioning that we can also extend the results to other pooling operations, including smoothed max pooling and attention-based pooling [45]. Due to limit of space, we include discussions in Appendix C.3 as well.
We justify the effectiveness of the proposed SONX and SONT algorithms for TPAUC Maximization in the regular learning setting and MIL setting [14, 45].
Table 2: Testing TPAUC on molecule datasets (top) and on MIL datasets (bottom). The two numbers in parentheses of the second line refers to the lower bound of TPR and the upper bound of FPR for evaluating TPAUC. The two numbers of each method refers to the mean TPAUC and its std.
moltox21 (t0)	molmuv (t1)	molpcba (t0)
Method	(0.6, 0.4)	(0.5, 0.5)	(0.6, 0.4)	(0.5, 0.5)	(0.6, 0.4)	(0.5, 0.5)
CE	0.067 (0.001)	0.208 (0.001)	0.161 (0.034)	0.469 (0.018)	0.095 (0.001)	0.264 (0.001)
AUC-SH	0.064 (0.008)	0.217 (0.014)	0.260 (0.130)	0.444 (0.128)	0.140 (0.003)	0.312 (0.003)
AUC-M	0.066 (0.009)	0.209 (0.01)	0.114 (0.079)	0.433 (0.053)	0.142 (0.009)	0.313 (0.003)
MB	0.067 (0.015)	0.215 (0.023)	0.173 (0.153)	0.426 (0.118)	0.095 (0.002)	0.262 (0.003)
AW-poly	0.064 (0.01)	0.206 (0.025)	0.172 (0.144)	0.393 (0.123)	0.110 (0.001)	0.281 (0.002)
SOTA-s	0.068 (0.018)	0.23 (0.021)	0.327 (0.164)	0.526 (0.122)	0.143 (0.001)	0.314 (0.002)
SONX	0.07 (0.035)	0.252 (0.025)	0.347 (0.175)	0.575 (0.122)	0.158 (0.006)	0.335 (0.006)
MUSK2	Fox
Method	(0.5, 0.5)	(0.3, 0.7)	(0.1, 0.9)	(0.5, 0.5)	(0.3, 0.7)	(0.1, 0.9)
AUC-M (att)	0.675 (0.1)	0.783 (0.067)	0.867 (0.036)	0.032 (0.03)	0.253 (0.098)	0.444 (0.118)
MIDAM (smx)	0.525 (0.2)	0.667 (0.149)	0.8 (0.097)	0.048 (0.059)	0.265 (0.119)	0.449 (0.113)
MIDAM (att)	0.6 (0.215)	0.717 (0.135)	0.819 (0.092)	0.016 (0.032)	0.249 (0.125)	0.509 (0.065)
SOTAs (att)	0.6 (0.267)	0.683 (0.178)	0.819 (0.097)	0.024 (0.032)	0.278 (0.059)	0.477 (0.046)
SONT (att)	0.7 (0.1)	0.8 (0.067)	0.867 (0.036)	0.12 (0.131)	0.343 (0.176)	0.578 (0.119)
Colon	Lung
Method	(0.5, 0.5)	(0.3, 0.7)	(0.1, 0.9)	(0.5, 0.5)	(0.3, 0.7)	(0.1, 0.9)
AUC-M (att)	0.576 (0.1)	0.739 (0.061)	0.803 (0.038)	0.32 (0.181)	0.609 (0.113)	0.744 (0.082)
MIDAM (smx)	0.646 (0.083)	0.787 (0.04)	0.863 (0.026)	0.43 (0.195)	0.68 (0.128)	0.824 (0.055)
MIDAM (att)	0.548 (0.253)	0.738 (0.149)	0.826 (0.102)	0.544 (0.261)	0.716 (0.189)	0.815 (0.129)
SOTAs (att)	0.772 (0.124)	0.862 (0.073)	0.911 (0.045)	0.539 (0.153)	0.745 (0.077)	0.841 (0.049)
SONT (att)	0.8 (0.166)	0.875 (0.099)	0.916 (0.065)	0.639 (0.137)	0.779 (0.041)	0.865 (0.028)
Baselines. For regular TPAUC maximization, we compare SONX with the following competitive methods: 1) Cross Entropy (CE) loss minimization; 2) AUC maximization with squared hinge loss (AUC-SH); 3) AUC maximization with min-max margin loss (AUC-M) [37]; 4) Mini-Batch based heuristic loss (MB) [16]; 5) Adhoc-Weighting based method with polynomial function (AW-poly) [35]; 5) a single-loop algorithm (SOTAs) for optimizing a smooth surrogate for TPAUC [44]. For MIL TPAUC maximization, we consider the following baselines: 1) AUC-M with attention-based pooling (AUC-M [att]); 2) SOTAs with attention-based pooling, which is a natural combination between advanced TPAUC optimization and MIL pooling technique; 3) the recently proposed provable multi-instance deep AUC maximization methods with stochastic smoothed-max pooling and attention-based pooling (MIDAM [smx] and MIDAM [att]) [45]. The first two baselines use naive mini-batch pooling for computing the loss function in AUC-M and SOTAs. We implement SONT for MIL TPAUC maximization with attention-based pooling, which is referred to as SONT (att).
Datasets. For regular TPAUC maximization, we use three molecule datasets as in [44], namely moltox21 (the No.0 target), molmuv (the No.1 target) and molpcba (the No.0 target) [29]. For MIL TPAUC maximization, we use four MIL datasets, including two tabular datasets MUSK2 and Fox, and two medical image datasets Colon and Lung. MUSK2 and Fox are two tabular datasets that have been widely adopted for MIL benchmark study [14]. Colon and Lung are two histopathology (medical image) datasets that have large image size  $(512\times 512)$  but local interests for classification [2]. For Colon dataset, the adenocarcinoma is regarded as positive label and benign is negative; for Lung dataset, we treat adenocarcinoma as positive and squamous cell carcinoma as negative  $^{3}$ . For both of the histopathology datasets, we uniformly randomly sample 100 positive and 1000 negative data for experiments. For all MIL datasets, we uniformly randomly split  $10\%$  as the testing and the remaining as the training and validation. The statistics for all used datasets are summarized in Table 3 and Table 4 in Appendix F.
Experiment Settings. For regular TPAUC maximization, we use the same setting as in [44]. The adopted backbone Graph Nueral Network (GNN) model is Graph Isomorphism Network (GIN), which has 5 mean-pooling layers with 64 number of hidden units and dropout rate 0.5 [30]. We utilize the sigmoid function for the final output layer to generate the prediction score, and set the surrogate loss  $\ell(\cdot)$  as squared hinge loss with a margin parameter. We follow the setups for model training and tuning exactly the same as the prior work [44]. Essentially, the model is trained by 60 epochs and the learning rate is decreased by 10-fold after every 20 epochs. The model is initialized as a pretrained model from CE loss on the training datasets. We fix the learning rate of SONX as 1e-2 and moving average parameter  $\tau$  as 0.9; tune the parameter  $\gamma$  in  $\{0, 1\mathrm{e} - 1, 1\mathrm{e} - 2, 1\mathrm{e} - 3\}$ , the parameter  $\alpha, \beta$  in  $\{0.1, 0.3, 0.5\}$  and fix the margin parameter of the surrogate loss  $\ell$  as 1.0, which cost the same
3Data available: https://www.kaggle.com/datasets/biplobdey/lung-and-colon-cancer
(a) molmuv (t1)
(b) molpcba (t0)
Figure 1: Training Curves of SONX (left two) and SONT (right two) for TPAUC maximization with different  $\gamma$ . The y-axis is the TPAUC (0.5, 0.5).
(c) MUSK2
(d) Fox
tuning effort as the other baselines. The weight decay is set as the same value (2e-4) with the other baselines. For baselines, we directly use the results reported in [44] since we use the same setting.