PaperlayerView exact revision on arXiv

Source-linked AI summary

ABO-Med: Accelerated Bilevel Optimization for Few-Shot Medical Image Classification

Ruoxuan Shi, Sheng Yang, Zhengxing Su, Xiaoyang Hou, Yating Liu

arXiv:2609.33176v1cs.CV

TL;DR

Because second-order bilevel methods can be costly and medical image classification often has limited labeled data, this paper develops ABO-Med, a first-order MAML-type framework with modality-aware MedRAug. Across several public medical datasets, ABO-Med outperformed prior baselines by around 1.99%–18.76%, while MedRAug improved average accuracy by around 2.20%–6.34%.

  • Problem

    Existing MAML bilevel optimization can require costly second-order computations, leaving the question of a nearly optimal first-order approach for medical few-shot learning.

    Source: Introduction · pp. 1–2

  • Method

    ABO-Med integrates first-order bilevel penalty methods into ANIL for medical few-shot classification and adds MedRAug, a modality-aware augmentation strategy tailored to medical imaging.

    Source: Introduction · p. 2

  • Results

    Across several public medical datasets, ABO-Med consistently outperformed prior state-of-the-art methods by around 1.99%–18.76%, while MedRAug raised average accuracy by around 2.20%–6.34%.

    Source: Result Analysis for Classification… · p. 6, Conclusion · p. 7

  • Takeaways & Limitations

    The paper supports ABO-Med as an efficient and effective approach to few-shot medical image classification.

    Source: Conclusion · p. 7

    The analysis assumes task losses satisfy standard conditions, including strong convexity in task variables and smoothness and Lipschitz properties.

    Source: Technical Assumptions and Definitions · p. 4

Abstract

from arXiv · show

In recent years, bilevel optimization has been widely used in a variety of machine learning tasks. However, prior bilevel optimization algorithms generally require the computation of second-order information, which limits their practical scalability. Only recently has a first-order paradigm for bilevel optimization been established, attaining near-optimal theoretical guarantees for solving bilevel optimization problems. In this paper, we propose ABO-Med, a scalable instantiation of this paradigm for few-shot learning, by incorporating it into the model-agnostic meta-learning (MAML) framework and tailoring it to medical image classification. We also introduce Medical Adaptive RandomAugment (MedRAug), a modality-aware augmentation strategy designed for medical images. Theoretically, ABO-Med establishes the optimality of MAML-type meta-learning approaches. Empirically, ABO-Med outperforms prior baselines on several public medical datasets, with gains of 1.99% to 18.76%, while MedRAug further improves the average accuracy by 2.20% to 6.34%. Additional cross-domain experiments, augmentation ablation studies, backbone ablation studies, and training efficiency analysis further validate the effectiveness and efficiency of the proposed method.

Introduction

Limited expert-annotated medical data motivates few-shot learning, while MAML-style adaptation can incur costly second-order computation. ABO-Med combines first-order bilevel optimization with modality-aware MedRAug; the paper reports benchmark gains and near-optimal convergence guarantees.

Introduction

  • Expert annotation is costly and time-consuming, limiting labeled medical samples; MAML's inner-loop differentiation adds Hessian computation and memory overhead, motivating nearly optimal first-order adaptation.The introduction poses this as whether a nearly optimal first-order MAML-type method can be developed.Source: Introduction · pp. 1–2

Contributions

  • ABO-Med is a scalable, flexible MAML-type framework for few-shot medical image classification with near-optimal convergence guarantees and no second-order derivative computations.It integrates first-order bilevel penalty methods into ANIL and incorporates MedRAug, a modality-aware strategy tailored to different medical imaging data types.Source: Contributions · p. 2, Introduction · p. 2
  • Average gains of 1.99%, 2.06%, and 18.76% on Pap-Smear, ISIC2018, and BreakHis40X, respectively, show ABO-Med outperforming strong baselines.The comparison is reported across multiple few-shot medical image classification benchmarks.Source: Contributions · p. 2
  • MedRAug further improves ABO-Med by 1.69%, 2.72%, and 6.01% on Pap-Smear, ISIC2018, and BreakHis40X, respectively.The paper also reports cross-domain, ablation, and efficiency studies verifying ABO-Med's effectiveness and efficiency.Source: Contributions · p. 2

Related Work

  • Prior bilevel methods include Hessian- and Jacobian-based hypergradient methods and first-order alternatives, while MAML variants and medical augmentation methods target adaptation cost and modality-specific distortions.MAML variants include FOMAML, Reptile, ANIL, ALFA, and Tra-MAML (Finn, Abbeel, and Levine 2017; Nichol, Achiam, and Schulman 2018; Raghu et al. 2019; Baik et al. 2020; Voon et al. 2025); prior augmentation work includes RandAugment and pathology-specific H&E-RandomAugment (Cubuk et al. 2020; Faryna, Van der Laak, and Litjens 2021).Source: Related Work · pp. 2–3

Preliminaries

The section defines notation, formulates few-shot meta-learning as a bilevel problem over support and query sets, and states assumptions on task losses used in the analysis.

Preliminaries

  • The paper uses partial gradients, the Euclidean norm, and asymptotic notation that hides absolute constants, optionally suppressing polylogarithmic factors.The notation O(·) hides only absolute constants independent of problem parameters; ˜O(·) additionally suppresses polylogarithmic factors, and Ω(·) denotes the corresponding lower-bound order.Source: Preliminaries · p. 3

Problem Formulation

  • Few-shot tasks are k-way, n-shot problems with support sets for adaptation and query sets for evaluation, framed as finite-sum bilevel optimization.Shared parameters ϕ define the feature extractor, while task-specific parameters w_i correspond to the last linear layer and minimize strongly convex support losses; the query objective is generally nonconvex in ϕ, with full-batch updates.Source: Problem Formulation · pp. 3–4

Technical Assumptions and Definitions

  • Task losses are assumed to have strongly convex, smooth inner objectives and Lipschitz, smooth, twice continuously differentiable outer objectives, with lower-bounded task meta-losses.The paper defines ℓ as the maximum of the listed smoothness constants and κ := ℓ/µ as the condition number.Source: Technical Assumptions and Definitions · p. 4

Method

ABO-Med trains a shared meta-parameter using a fully first-order penalized bilevel method, then adapts and evaluates it on unseen tasks. Its theory gives a meta-training stationarity guarantee and a separate test-time adaptation guarantee.

ABO-Med

  • ABO-Med samples task batches, averages approximate hypergradients to update a shared meta-parameter, then adapts it on unseen-task support sets with AGD and reports average query accuracy.Meta-training tasks are split into support and query sets; meta-testing fixes and transfers the learned meta-parameter to sampled unseen tasks.Source: ABO-Med · p. 4

Fully First-Order MAML Method

  • The penalized bilevel formulation yields a hypergradient using only first-order support- and query-loss gradients, avoiding the Hessian- or Jacobian-related computations of the original objective.For fixed ϕ, gradient descent on the penalized and support losses approximates the auxiliary variables ˜w∗λ(ϕ) and ˜w∗(ϕ), respectively.Source: Fully First-Order MAML Method · p. 5

Theoretical Results

  • With α = 1/Ls, τ = 1/(2λLs), and K = O(κ log(λℓ/µ)), meta-training reaches an ϵ-first-order stationary point within T = O(ℓκ^4ϵ^-2 log(ℓκ/ϵ)) first-order oracle calls.The guarantee assumes Assumption 1; choosing K′ = e O(√κ) makes AGD produce a sufficiently accurate approximation of each test-task minimizer.Source: Theoretical Results · p. 5

Experiments

Experiments show that ABO-Med improves few-shot classification accuracy over competing methods, while evaluation also examines MedRAug, training efficiency, and cross-domain generalization.

Experiments

  • The experiments cover classification performance, augmentation and backbone ablations, training efficiency, and cross-domain generalization.Additional backbone results are provided in Appendix A.5.Source: Experiments · p. 5

Dataset Description

  • The evaluation uses Pap-Smear (917 images, 7 classes), ISIC2018 (10,015 images, 7 classes), and BreaKHis40X (1,820 images, 8 classes), with BCHI (922 images, 3 classes) as the cross-domain target.BCHI is used to evaluate generalization under domain shift.Source: Dataset Description · p. 5

Medical Adaptive RandomAugment

  • MedRAug adapts RandomAugment with modality-specific operations and magnitude ranges to increase diversity while preserving clinically relevant medical-image structures.Its ablation compares it with baseline augmentation, CutMix, Mixup, and Cutout.Source: Medical Adaptive RandomAugment · pp. 5–6

Result Analysis for Classification Performance

  • 1.99%, 2.06%, and 18.76%: ABO-Med improves over the best competing method on Pap-Smear, ISIC2018, and BreakHis40X, respectively.With MedRAug, average accuracy improves by 2.20%, 3.66%, and 6.34%, achieving the best results across all datasets and shot settings.Source: Table 1 · p. 6, Result Analysis for Classification… · p. 6

Training Time Analysis

  • Approximately 19.0%, 73.1%, and 72.8%: ABO-Med reduces average training time versus the best competing method on Pap-Smear, ISIC2018, and BreakHis40X, respectively.Training time was measured per 100 tasks under the same 3-way 1-, 3-, 5-, 7-, and 10-shot settings, CNN4 backbone, and baseline augmentation.Source: Table 2 · p. 6, Training Time Analysis · p. 7

Cross-domain Analysis

  • 2.16% and 10.64%: ABO-Med improves over the strongest baseline in near-domain and far-domain cross-domain settings, respectively.The experiments meta-train on BreakHis40X or Pap-Smear and meta-test on BCHI; the results suggest ABO-Med learns more stable, transferable representations under domain shift.Source: Table 3 · p. 7, Cross-domain Analysis · p. 7

Conclusion

ABO-Med is a scalable MAML-type framework for few-shot medical image classification, with convergence and complexity guarantees under a fully first-order bilevel formulation. Experiments show improvements over prior methods, and MedRAug further improves average accuracy.

Conclusion

  • 1.99% to 18.76% improvements over previous state-of-the-art methods were reported across several public medical datasets.The conclusion describes these gains as consistent across datasets.Source: Conclusion · p. 7
  • 2.20% to 6.34% average-accuracy improvements were reported with MedRAug, whose effectiveness was validated through augmentation ablation studies.Source: Conclusion · p. 7
  • ABO-Med is a scalable and flexible MAML-type framework with convergence guarantees and complexity results under a fully first-order bilevel formulation.The paper presents it as an efficient and effective solution for few-shot medical image classification, where stable adaptation and computational scalability are critical.Source: Conclusion · p. 7

A Appendix

ABO-Med’s many hyperparameters make careful tuning necessary and may hinder practical use when validation is costly; future work will pursue adaptive optimization and automated hyperparameter workflows.

Limitations and Future Work

  • ABO-Med’s many hyperparameters require careful tuning across tasks and datasets, increasing implementation complexity and potentially limiting practicality when validation is costly.Source: A Appendix · p. 11
  • Future work will develop adaptive bilevel optimization methods and explore agent-based workflows for selecting and adjusting key hyperparameters.The cited works include Huang, Li, and Gao 2021; Fan et al. 2023; Antonakopoulos et al. 2025; Yang et al. 2024; Zhai, Yan, and Zhang 2025 for adaptive methods, and Rao et al. 2025; Ma et al. 2026; Rao et al. 2026 for automated workflows.Source: A Appendix · p. 11

A.1 Extend Related Work

This section extends the related work with optimization-based meta-learning methods that complement the main text. It describes their adaptation or optimization strategies and notes that they were excluded from the main empirical comparison because they did not outperform the representative baselines and state-of-the-art methods considered.

A.1 Extend Related Work

  • Meta-SGD learns parameter-wise learning rates for fast adaptation, while MAML++ uses optimization refinements to improve MAML’s training stability (Li et al. 2017; Antoniou, Edwards, and Storkey 2018).Both methods extend or refine MAML.Source: §A.1 · p. 11
  • iMAML uses implicit differentiation at the inner-level solution to compute meta-gradients, avoiding backpropagation through the inner optimization path and reducing memory overhead (Rajeswaran et al. 2019).Source: §A.1 · p. 11
  • FO-B-MAML presents a fully first-order bilevel approach that regularizes the inner problem and approximates hypergradients via finite differences (Chayti and Jaggi 2024).Source: §A.1 · p. 11
  • The authors exclude these methods from the main empirical comparison because they do not outperform the representative baselines and state-of-the-art methods considered in their experiments.Source: §A.1 · p. 11

A.2 Proof for Theoretical Results

Under Assumption 1, Algorithm 1’s meta-training stage reaches an ϵ-first-order stationary point within the stated first-order oracle complexity, which matches an Ω(ϵ^-2) lower bound.

A.2 Proof for Theoretical Results

  • For λ≥2Ld/µ, Proposition 5 gives O(ℓκ^3)-gradient-Lipschitzness for each penalized task objective, following Lemma 3.1 in (Kwon et al. 2023) and Lemma 4.1 in (Chen, Ma, and Zhang 2025).Source: §A.2 · p. 11
  • Lemma 7 bounds the squared gap between penalized and support-loss minimizers by O(κ^2/λ^2) when λ≥2Ld/µ and Cd≤ℓ.Source: §A.2 · p. 11
  • Lemma 8 provides an inexact-descent bound for an L-gradient-Lipschitz objective when the update uses η≤1/(2L).Source: §A.2 · p. 11
  • O(ℓκ^4ϵ^-2 log(ℓκ/ϵ)) first-order oracle calls suffice for Algorithm 1’s meta-training stage to find an ϵ-first-order stationary point of L(ϕ).Under Assumption 1, the theorem sets α=1/Ls, τ=1/(2λLs), and K=O(κ log(λℓ/µ)).Source: §A.2 · p. 12
  • Meta-testing separately chooses K′=eO(√κ) for AGD to approximate each test-task minimizer, and K′ is excluded from meta-training oracle complexity.Source: §A.2 · p. 12
  • The bound is near-optimal because it matches the Ω(ϵ^-2) lower complexity bound.Source: §A.2 · p. 14

A.3 Dataset Split

The evaluation used separate meta-training and meta-testing classes to assess generalization to unseen classes, with BCHI reserved as the cross-domain target.

A.3 Dataset Split

  • Pap-Smear and ISIC2018 used 4 meta-training and 3 meta-testing classes, BreaKHis40X used 5 and 3, and BCHI was only the cross-domain target with all 3 grades meta-tested.Source: §A.3 · p. 14

A.4 Experiment Details

Experiments used a shared comparison protocol with a fixed CNN4 backbone and consistent few-shot settings. ABO-Med hyperparameters were selected by validation, with the same configuration used for MedRAug experiments.

A.4 Experiment Details

  • All methods shared the backbone, few-shot setting, and data-augmentation strategy for fair comparison, using CNN4 and 3-way 3-, 5-, and 10-shot evaluation.CNN4 used 84×84 inputs and four 3×3 convolutional layers with 64 filters each, followed by batch normalization and ReLU.Source: §A.4 · p. 14
  • Meta-training used task batches of 32; during testing, the best weights from meta-training were evaluated on 600 meta-testing tasks, with 30 adaptation steps per task.The baseline augmentation matched Tra-MAML, including random resized crop, image jitter, and random vertical and horizontal flips.Source: §A.4 · p. 14
  • ABO-Med hyperparameters were chosen by grid search for best validation performance, and MedRAug used the same settings to isolate augmentation effects.Search ranges covered η, α, τ, λ, K, K′, β, and θ; final settings are identified as reported in Table 4, whose values are not included in the supplied passage.Source: §A.4 · p. 14

A.5 Ablation Study on Backbones

Across medical few-shot backbone ablations, lightweight CNN4 achieved the best overall performance, while deeper ResNets and ViTFeatsTiny provided no further improvements.

A.5 Ablation Study on Backbones

  • Medical few-shot tasks rely on subtle morphological or texture differences, so greater backbone capacity may overfit or require more data and regularization.The authors contrast these fine-grained tasks with natural-image few-shot benchmarks, where larger semantic differences and visual diversity can benefit high-capacity backbones; CNN4 balanced capacity and complexity and performed more stably across datasets and shot settings.Source: §A.5 · pp. 14–15
  • CNN4 achieved the best overall performance despite having the fewest trainable parameters; deeper ResNets and ViTFeatsTiny did not improve results.The comparison used the same baseline augmentation settings on Pap-Smear, ISIC2018, and BreakHis40X under 3-way 3-, 5-, and 10-shot settings.Source: §A.5 · p. 14, Table 5 · p. 15

A.6 Ablation Study on Data Augmentation … Data Augmentation Strategy

MedRAug outperforms the compared augmentation strategies across datasets and shot settings, while the appendix explains how augmentation choices are tailored to medical-image modalities and how conventional methods can alter diagnostic patterns.

Details of Dataset-specific Augmentation Spaces

  • Augmentation spaces are tailored by dataset: Pap-smear uses moderate variation while preserving cell morphology, ISIC2018 targets lesion, angle, and illumination variation, and BreakHis40X preserves tissue structure while simulating staining and scanning variation.The appendix summarizes the selected operations and their design rationale for the three datasets.Source: Details of Dataset-specific… · p. 15

Experimental results on ablation study

  • The ablation changes only augmentation strategy, holds other settings fixed, uses the same backbone for every method, and compares baseline, MedRAug, CutMix, Mixup, and Cutout.The compared results are reported in Table 9 as mean test accuracy with 95% confidence intervals.Source: Experimental results on ablation study · p. 15, Table 9 · p. 17
  • CutMix, Mixup, and Cutout help in some BreakHis40X settings but generally degrade Pap-Smear and ISIC2018 performance, whereas MedRAug is consistently best.The authors suggest that direct image mixing or random masking can disrupt important medical visual patterns.Source: Experimental results on ablation study · p. 15

Data Augmentation Strategy

  • Cutout randomly masks a rectangular image region to discourage reliance on a few discriminative regions, but excessive occlusion can remove diagnostically important structures.The method is presented as a way to improve robustness and generalization, especially with limited data.Source: Data Augmentation Strategy · p. 16
  • Mixup forms virtual training samples by linearly combining two inputs and labels, encouraging smoother decision boundaries, though strong mixing may weaken diagnostic structures.Its mixing coefficient is usually sampled from a Beta distribution.Source: Data Augmentation Strategy · p. 16
  • CutMix replaces a rectangular region with a patch from another image and mixes labels according to replaced area, preserving more local structure while adding cross-class features.The described method aims to improve robustness to local perturbations and encourage discriminative representations.Source: Data Augmentation Strategy · pp. 16–17

Experimental results on ablation study

  • MedRAug achieves the best performance across all datasets and shot settings, improving average accuracy over baseline augmentation by 1.69% on Pap-Smear, 2.72% on ISIC2018, and 6.01% on BreakHis40X.The reported improvements are average accuracy gains over the baseline augmentation strategy.Source: Experimental results on ablation study · p. 15, Table 9 · p. 17

Related papers

Loading 2609.33176v1…