BalGrad: Balanced Gradient Optimization in VL Models

• BalGrad is a framework for vision-language models that addresses dominant modality bias by realigning gradient magnitudes and directions.
• It employs inter-modality gradient reweighting and inter-task gradient projection to enable balanced convergence between image and text modalities.
• Experimental results demonstrate improved performance under missing or noisy modalities compared to static loss weighting approaches.

BalGrad is a framework for vision-language (VL) models proposed to mitigate “dominant modality bias,” defined as a setting “where a specific modality disproportionately influences the model’s outcomes” and that “behaves particularly detrimentally when one modality is impaired,” such as under missing or noisy image or text inputs. Its central claim is that “unaligned gradients or differences in gradient magnitudes prevent balanced convergence of the loss,” so the method intervenes directly at the gradient level through inter-modality gradient reweighting and inter-task gradient projection rather than through static loss weighting alone (Kwon et al., 18 Mar 2025).

1. Dominant modality bias and task setting

In the BalGrad formulation, dominant modality bias denotes the tendency of a multimodal predictor to rely disproportionately on one modality for its decisions. The empirical manifestation is task-dependent. On Hateful Memes, text leads; on UPMC Food-101, images lead; on MM-IMDb, there is no clear dominance. When the dominant modality is impaired, performance drops sharply, which exposes over-reliance rather than genuinely balanced multimodal inference. The training dynamics also exhibit asymmetry: the loss of the dominant modality decreases rapidly, while the weak modality decreases slowly, indicating under-exploration of the weaker modality during optimization (Kwon et al., 18 Mar 2025).

The target task loss is written as $$\mathcal{L}(\theta_v, \theta_l, \theta_{\mathcal{T}})$$, where $\theta_v$ and $\theta_l$ are the parameters of image and text embedding layers, and $\theta_{\mathcal{T}}$ denotes the parameters of the classifier $f_{\mathcal{T}}(\cdot)$. The method decomposes the target-loss gradient into modality-specific components, denoted by $\mathcal{G}^\tau = \{g_l, g_v, g_{\mathcal{T}}\}$. This decomposition is not an auxiliary diagnostic only; it is the basis for defining when convergence is balanced and for identifying when one modality is driving optimization disproportionately.

BalGrad was developed for bi-modal VL settings in which the two modalities can be separately probed during training. To inspect the training dynamics, the framework attaches linear heads $f_v(\cdot)$ and $f_l(\cdot)$ on top of the image and text embedding layers, producing unimodal probabilities $p_i^v$ and $p_i^l$ and modality-specific target objectives $\theta_v$0 and $\theta_v$1. This makes the imbalance observable at the loss and gradient levels rather than only at the level of aggregate task metrics.

2. Gradient analysis and balanced convergence

BalGrad’s theoretical motivation is expressed through two propositions. The first states that the interaction between modality-specific gradients governs how much the overall loss is reduced. In the authors’ wording, “the cross term $\theta_v$2 captures the interaction between the visual and language gradients and the magnitudes and directions of each gradient $\theta_v$3 and $\theta_v$4 governs how much the overall loss is reduced.” The immediate consequence is that if these gradients are unaligned, or if their magnitudes differ substantially, “the loss reduction will not be balanced,” and the dominant modality drives larger decreases while the weak modality remains under-optimized (Kwon et al., 18 Mar 2025).

The second proposition concerns interference between the target loss $\theta_v$5 and a balance loss based on KL divergence. Let $\theta_v$6 denote gradients from the target loss and $\theta_v$7 denote gradients from the KL loss. Under gradient descent with small step size $\theta_v$8,

$\theta_v$9

If $\theta_l$0, then the target and KL gradients conflict, reducing the effective loss reduction. In BalGrad, this is the formal justification for introducing a projection step rather than simply adding the KL term to the target objective.

The “see-saw modality balance” intuition follows directly from these propositions. When gradient magnitudes differ or directions conflict, the optimization trajectory tips toward one modality. Balanced convergence therefore requires both magnitude correction and directional alignment. BalGrad implements the former through reweighting and the latter through projection.

3. Mechanism of BalGrad

BalGrad augments the target task objective with mutual learning via bidirectional KL divergence between unimodal predictions. The mutual Kullback–Leibler terms are defined as

$\theta_l$1

and symmetrically $\theta_l$2, with corresponding gradients $\theta_l$3 and $\theta_l$4. These terms provide a balance signal, but BalGrad does not treat them uniformly across training. Instead, it reweights them according to the current unimodal target losses:

$\theta_l$5

The more converged modality receives a lower KL-gradient weight, while the weaker modality receives a stronger alignment signal. The aggregated KL gradient is then scheduled over time as

$\theta_l$6

with initial weighting factor $\theta_l$7 (Kwon et al., 18 Mar 2025).

The second component is inter-task gradient projection. BalGrad first computes the dot product $\theta_l$8. If the result is nonnegative, the target gradient is used as-is. If the result is negative, indicating conflict, BalGrad projects the target gradient onto the orthogonal complement of $\theta_l$9:

$\theta_{\mathcal{T}}$0

This preserves target learning while removing the component that opposes modality balancing.

Operationally, the training loop is described as “see” and “sew.” The “see” stage computes $\theta_{\mathcal{T}}$1, $\theta_{\mathcal{T}}$2, $\theta_{\mathcal{T}}$3, $\theta_{\mathcal{T}}$4, and $\theta_{\mathcal{T}}$5, then reweights and aggregates the KL gradients into $\theta_{\mathcal{T}}$6. The “sew” stage applies the conflict test, replaces $\theta_{\mathcal{T}}$7 by $\theta_{\mathcal{T}}$8 when necessary, updates the final classifier $\theta_{\mathcal{T}}$9 with $f_{\mathcal{T}}(\cdot)$0, and updates the image and text embedding layers jointly with $f_{\mathcal{T}}(\cdot)$1, $f_{\mathcal{T}}(\cdot)$2, $f_{\mathcal{T}}(\cdot)$3, $f_{\mathcal{T}}(\cdot)$4, and $f_{\mathcal{T}}(\cdot)$5. The method therefore intervenes in both fused-task optimization and unimodal representation learning.

4. Experimental behavior under modality impairment

The main experiments use ViT (vit-base) for vision and BERT (bert-base-uncased) for language, with late concatenation. Encoders are frozen and only the embedding and classifier layers are trained. Each encoder produces a 768-dimensional embedding, and concatenation yields a 1568-dimensional vector. Images are resized to $f_{\mathcal{T}}(\cdot)$6 with random horizontal flip; maximum text sequence lengths are 1024 for MM-IMDb, 512 for UPMC Food-101, and 128 for Hateful Memes. Optimization uses Adam with momentum $f_{\mathcal{T}}(\cdot)$7, training for 20 epochs with batch size 128 (Kwon et al., 18 Mar 2025).

On UPMC Food-101, which contains 90,840 image-text pairs across 101 food classes and uses accuracy as the metric, BalGrad improves both full-condition and impaired-modality performance. Under the full setting it reports 80.32, compared with 76.01 for the baseline. Under missing image it reports 25.49, compared with 22.60 for AGM and 12.99 for the baseline. Under missing text it reports 65.03, compared with 63.52 for the baseline. The average under missing conditions is 45.26, the highest among the compared methods, and $f_{\mathcal{T}}(\cdot)$8Gap under missing conditions is 39.54, the smallest. Under noisy conditions, the average is 67.06, again the highest (Kwon et al., 18 Mar 2025).

On Hateful Memes, with 8,500 train, 1,000 validation, and 500 test examples and AUROC as the metric, BalGrad reports 67.35 in the full setting, 65.86 under missing image, and 57.58 under missing text. Its average under missing conditions is 61.72 and $f_{\mathcal{T}}(\cdot)$9Gap under missing conditions is 8.28, both the best among the reported methods. Under noisy conditions, the average is 65.69 and $\mathcal{G}^\tau = \{g_l, g_v, g_{\mathcal{T}}\}$0Gap is 0.18. The experiments also note negative transfer for OGM-GE and AGM on this dataset: image-missing performance exceeds full-modality performance, which the authors interpret as a sign that “adding the image modality results in a decrease in performance compared to using text alone.”

On MM-IMDb, a multi-label movie genre classification benchmark with 15,552 train, 2,608 validation, and 7,799 test samples and F1-Macro as the metric, the pattern is different. BalGrad reports 43.19 in the full setting, slightly lower than the baseline at 44.09; 18.81 under missing image, near the baseline at 18.85; and 17.47 under missing text, second highest against a baseline of 18.40. Even in this less biased regime, its noisy-condition average is 39.78, the highest among the compared methods, and $\mathcal{G}^\tau = \{g_l, g_v, g_{\mathcal{T}}\}$1Gap under noisy conditions remains low at 4.04. This suggests that BalGrad is most beneficial when dominant modality bias is pronounced, while remaining comparatively stable when no clear dominance exists.

5. Ablations, generalization, and diagnostic analyses

BalGrad’s component analysis separates the effects of reweighting and projection. Reweighting alone improves averages and reduces gaps but can induce negative transfer; on Hateful Memes, image-missing performance can exceed full performance, and on MM-IMDb the model remains overly reliant on text. Projection alone reduces conflict and negative transfer but is less effective than the full method. The combined system performs best overall, and the authors summarize the mechanism as follows: “By aligning the gradient of the target loss with the KL loss term, we reduce reliance on any single modality” (Kwon et al., 18 Mar 2025).

The diagnostic evidence is explicitly gradient-based. Loss curves per modality show that reweighting yields “faster and more stable convergence of loss for each modality.” Histograms of cosine similarity between $\mathcal{G}^\tau = \{g_l, g_v, g_{\mathcal{T}}\}$2 and $\mathcal{G}^\tau = \{g_l, g_v, g_{\mathcal{T}}\}$3 show that negative similarity is prevalent without projection, while projection shifts the mean cosine to positive values. The conflict fraction and $\mathcal{G}^\tau = \{g_l, g_v, g_{\mathcal{T}}\}$4Gap decrease in parallel: on UPMC Food-101, the conflict fraction decreases from 0.66 to 0.36 and $\mathcal{G}^\tau = \{g_l, g_v, g_{\mathcal{T}}\}$5Gap from 43.27 to 39.54; on Hateful Memes, from 0.78 to 0.32 and from 10.21 to 8.28; on MM-IMDb, from 0.28 to 0.26 and from 4.21 to 4.04. The asymmetry of these changes is consistent with the dataset-level bias patterns: large gains where modality bias is strong, marginal changes where bias is mild.

Generalization experiments extend beyond the main ViT+BERT setting. On BLIP, a decoder-based VL model, BalGrad achieves the lowest $\mathcal{G}^\tau = \{g_l, g_v, g_{\mathcal{T}}\}$6Gap across datasets. Across fusion mechanisms, including addition and attention, it consistently yields the smallest $\mathcal{G}^\tau = \{g_l, g_v, g_{\mathcal{T}}\}$7Gap. Across backbones, including ResNet-50 + DistilBERT and CLIP, it again maintains the smallest $\mathcal{G}^\tau = \{g_l, g_v, g_{\mathcal{T}}\}$8Gap. The MM-IMDb case is especially informative here: under weaker backbones it exhibits text bias, and BalGrad mitigates this, which supports the claim that the method addresses optimization imbalance rather than only a peculiarity of one architecture.

Additional datasets reinforce the same pattern. On Memotion, measured by weighted F1, BalGrad reports Avg 56.13 and $\mathcal{G}^\tau = \{g_l, g_v, g_{\mathcal{T}}\}$9Gap 3.35, both the best across methods, with missing-text performance of 52.78. On CUB-200-2011, measured by accuracy, it reports missing image 45.47, more than 4%p over AGM at 41.24, Avg 54.10, and $f_v(\cdot)$0Gap 8.63, the smallest among the reported methods. These results indicate that the framework transfers beyond the three primary benchmarks and remains effective in fine-grained and multimodal affective settings.

6. Broader balanced-gradient landscape and limitations

BalGrad, in its strictest sense, denotes the VL method based on inter-modality gradient reweighting and inter-task gradient projection. A plausible broader taxonomy is that the term also sits within a larger class of balanced-gradient procedures that rebalance, normalize, project, or coordinate gradients in order to control conflict, scale disparity, or optimization drift. Within multi-task learning, "Feed m Birds with One Scone: Accelerating Multi-task Gradient Balancing via Bi-level Optimization" interprets “BalGrad” as the family of balanced-gradient methods typified by MGDA, CAGrad, and Nash-MTL, and proposes MARIGOLD as a bi-level zeroth-order acceleration framework that reduces per-iteration cost from $f_v(\cdot)$1 to $f_v(\cdot)$2 (Chen et al., 8 Mar 2026). In continuous-time multi-objective optimization, "Multiobjective Balanced Gradient Flow" defines the Multiobjective Balanced Gradient Flow through projection onto the convex hull of normalized gradients and proves convergence to weak Pareto points in the convex case, with rates of $f_v(\cdot)$3 and $f_v(\cdot)$4 in convex and non-convex settings respectively (Yin, 3 Aug 2025).

The same balancing motif appears in other optimization domains. GraB and CD-GraB use balanced prefix sums of per-example gradients to construct data permutations that provably outperform random reshuffling, with CD-GraB extending the idea to distributed training and establishing linear-in-$f_v(\cdot)$5 speedups relative to centralized GraB (Lu et al., 2022, Cooper et al., 2023). Fed-GraB introduces a Self-adjusting Gradient Balancer and a Direct Prior Analyzer for federated long-tailed learning, reweighting class-wise positive and negative gradients under privacy constraints (Xiao et al., 2023). In dense-prediction unsupervised domain adaptation, Gradient-based class weighting derives class weights from per-class loss gradients and improves tail-class recall without explicit target-prior estimation (Alcover-Couso et al., 2024). In face recognition fairness, the Gradient Attention Balance Network aligns gradient attention maps across races and uses GAM-guided masking to mitigate racial bias (Huang et al., 2023). These works do not define the same algorithm, but they share a common design principle: imbalance is treated as a gradient-geometry problem.

BalGrad also has explicit limitations. It is developed for bi-modal settings, and “extending this approach to multimodal models with more than two modalities presents additional challenges,” because pairwise gradient relations increase training cost and computational demands. The main experiments use frozen encoders under linear probing, and although the method generalizes across backbones and fusion schemes, performance under full fine-tuning and large-scale pretraining may require careful hyperparameter tuning. The framework further relies on unimodal heads to measure $f_v(\cdot)$6 and $f_v(\cdot)$7, so its most direct application is to models in which modality-specific supervision and gradient extraction are architecturally accessible. Within those limits, BalGrad’s significance lies in making modality balance an explicit gradient-level optimization objective rather than a by-product of multimodal fusion (Kwon et al., 18 Mar 2025).