PCGrad and Dynamic Weight Average Optimization for Balancing Contrastive Semantic Loss and Multi-Class Intent Cross-Entropy
Multi-Task Gradient Topography in Search Representation Learning
Deploying a modern, production-grade keyword cluster tool requires vector representations that simultaneously encode two distinct feature spaces: fine-grained semantic proximity and functional search intent. While transformer-based bi-encoders excel at capturing lexical and conceptual overlap through metric learning, standard single-task architectures fail to isolate search engine results page (SERP) intent modalities, often collapsing transactional queries into informational thematic clusters due to surface-level token overlap.
To resolve this, joint training paradigms optimize shared backbone parameters across two disparate loss functions: a contrastive semantic loss (such as InfoNCE or Multiple Negatives Ranking Loss) and a categorical cross-entropy loss over discrete intent labels. However, naive scalar addition of these losses introduces acute multi-task learning (MTL) pathologies, predominantly gradient interference, destructive negative transfer, and magnitude domination. Achieving Pareto optimality across this dual-objective manifold requires an interconnected optimization framework combining Projecting Conflicting Gradients (PCGrad) with Dynamic Weight Average (DWA) loss scheduling, expanding upon the foundational architecture outlined in Multi-Task Bi-Encoder Fine-Tuning for Joint Search Intent and Semantic Similarity Scoring.
The Multi-Task Objective Formulation
Let $mathbf{Theta}_{enc}$ denote the shared parameters of a bi-encoder transformer backbone (e.g., ModernBERT or RoBERTa), $mathbf{Theta}_{proj}$ represent the parameter set for the contrastive projection head, and $mathbf{Theta}_{cls}$ denote the parameter set of the linear intent classification head. Given a mini-batch $mathcal{B}$ containing query strings $x_i$, positive paired targets $x_i^+$, and multi-class intent labels $y_i in {1, dots, C}$ (where $C$ denotes intent classes: Informational, Commercial, Transactional, Navigational):
Contrastive Semantic Loss ($mathcal{L}_{cont}$)
The contrastive objective enforces high cosine similarity between semantically congruent query representations while repelling in-batch negatives. Utilizing an InfoNCE framework with temperature $tau$:
$$mathcal{L}_{cont}(mathbf{Theta}_{enc}, mathbf{Theta}_{proj}) = -frac{1}{|mathcal{B}|} sum_{i in mathcal{B}} log frac{exp(text{sim}(mathbf{z}_i, mathbf{z}_i^+) / tau)}{sum_{j in mathcal{B}} exp(text{sim}(mathbf{z}_i, mathbf{z}_j^+) / tau)}$$
where $mathbf{z}_i = mathbf{Theta}_{proj}(f(x_i; mathbf{Theta}_{enc}))$ is the normalized projected embedding, and $text{sim}(mathbf{u}, mathbf{v}) = frac{mathbf{u} cdot mathbf{v}}{|mathbf{u}|_2 |mathbf{v}|_2}$.
Intent Classification Cross-Entropy Loss ($mathcal{L}_{intent}$)
Concurrently, pooled representation vectors $mathbf{h}_i = f(x_i; mathbf{Theta}_{enc})$ pass through the classification head to optimize standard multi-class categorical cross-entropy:
$$mathcal{L}_{intent}(mathbf{Theta}_{enc}, mathbf{Theta}_{cls}) = -frac{1}{|mathcal{B}|} sum_{i in mathcal{B}} sum_{c=1}^C y_{i,c} log hat{y}_{i,c}$$
where $hat{mathbf{y}}_i = text{softmax}(mathbf{Theta}_{cls} mathbf{h}_i)$.
Pathologies of the Naive Composite Gradient
In standard gradient descent, the joint loss is expressed as $mathcal{L}_{total} = mathcal{L}_{cont} + mathcal{L}_{intent}$. Let $mathbf{g}_{cont} = nabla_{mathbf{Theta}_{enc}} mathcal{L}_{cont}$ and $mathbf{g}_{intent} = nabla_{mathbf{Theta}_{enc}} mathcal{L}_{intent}$ denote the gradient vectors with respect to the shared encoder backbone. Two critical failures emerge:
- Gradient Directional Conflict: When the inner product $langle mathbf{g}_{cont}, mathbf{g}_{intent} rangle < 0$, the update vectors pull shared attention weights in opposing directions across parameter space, degrading semantic clustering quality to satisfy discrete classification boundaries or vice versa.
- Magnitude Asymmetry: $mathcal{L}_{cont}$ calculations over large in-batch negatives typically manifest gradient norms $|mathbf{g}_{cont}|$ an order of magnitude smaller than unnormalized logit gradients $|mathbf{g}_{intent}|$, leading to intent dominance and degraded geometric embedding spaces.
Dynamic Weight Average (DWA) Optimization
To dynamically counteract magnitude drift and training-velocity discrepancies without relying on computationally prohibitive Hessian computations (as seen in GradNorm), Dynamic Weight Average (DWA) regulates task weights at the epoch or iteration boundary based on the relative rate of loss change.
Let $mathcal{L}_k(t)$ represent the smoothed average loss value for task $k in {cont, intent}$ at training step $t$. The relative descent indicator $w’_k(t)$ is formulated as:
$$r_k(t) = frac{mathcal{L}_k(t)}{mathcal{L}_k(t-1)}$$
The operational loss weight $lambda_k(t)$ is subsequently computed using a softmin-normalized temperature mapping:
$$lambda_k(t) = frac{K cdot exp(r_k(t-1) / T)}{sum_{j=1}^K exp(r_j(t-1) / T)}$$
where $K=2$ is the task cardinality and $T > 0$ is a temperature hyperparameter. As $T to infty$, weights converge toward uniformity ($lambda_k to 1$); lower values of $T$ aggressively prioritize the task demonstrating a slower convergence rate (higher $r_k$), preventing the intent classification head from prematurely plateauing while the contrastive projection continues to learn intricate semantic topologies.
Projecting Conflicting Gradients (PCGrad) Mechanics
While DWA remediates macroscopic scale imbalances over epoch intervals, it cannot mitigate batch-level directional opposition. Projecting Conflicting Gradients (PCGrad) provides microscopic orthogonalization to prevent destructive gradient interference on shared parameters $mathbf{Theta}_{enc}$.
During backward propagation, the shared layer gradient updates are modified on a per-step basis:
- Compute the unprojected task-specific gradients:
$$mathbf{g}_1 = nabla_{mathbf{Theta}_{enc}} (lambda_{cont} mathcal{L}_{cont}), quad mathbf{g}_2 = nabla_{mathbf{Theta}_{enc}} (lambda_{intent} mathcal{L}_{intent})$$ - Assess the directional cosine:
$$cos phi = frac{mathbf{g}_1 cdot mathbf{g}_2}{|mathbf{g}_1|_2 |mathbf{g}_2|_2}$$ - If $mathbf{g}_1 cdot mathbf{g}_2 < 0$, project each gradient vector onto the normal plane of the opposing task gradient:
$$mathbf{g}_1^{proj} = mathbf{g}_1 – frac{mathbf{g}_1 cdot mathbf{g}_2}{|mathbf{g}_2|_2^2} mathbf{g}_2$$
$$mathbf{g}_2^{proj} = mathbf{g}_2 – frac{mathbf{g}_2 cdot mathbf{g}_1}{|mathbf{g}_1|_2^2} mathbf{g}_1$$ - Apply symmetric summation: If $mathbf{g}_1 cdot mathbf{g}_2 ge 0$, keep original gradients ($mathbf{g}_k^{proj} = mathbf{g}_k$). The parameter update is governed by:
$$mathbf{g}_{final} = mathbf{g}_1^{proj} + mathbf{g}_2^{proj}$$
This projection eliminates the negative projection component of $mathbf{g}_i$ that works directly against the descent trajectory of $mathbf{g}_j$, preserving directional consensus while maintaining the non-conflicting components of each loss surface.
Unified Algorithmic Pipeline for Keyword Embeddings
The combined optimization architecture executes within the forward-backward loop of the training framework. The integration of DWA with PCGrad within a multi-task bi-encoder training step operates as follows:
Algorithm Execution Sequence
- Forward Pass: Propagate tokenized batch sequences through the shared transformer backbone $mathbf{Theta}_{enc}$. Pass representation states to $mathbf{Theta}_{proj}$ and $mathbf{Theta}_{cls}$ to output scalar similarities and class logits simultaneously.
- Weight Dynamic Evaluation: Update loss ratios $r_{cont}$ and $r_{intent}$ using windowed moving averages. Calculate dynamic weights $lambda_{cont}(t)$ and $lambda_{intent}(t)$ via temperature-scaled softmin.
- Independent Backward Graph Traversals:
- Execute
loss_cont.backward(retain_graph=True)to populate $mathbf{g}_{cont}$ across $mathbf{Theta}_{enc}$. - Store $mathbf{g}_{cont}$ in parameter-level temporal tensors and zero the shared parameter gradients.
- Execute
loss_intent.backward()to populate $mathbf{g}_{intent}$ across $mathbf{Theta}_{enc}$.
- Execute
- Orthogonal Projection: Traverse shared parameter tensors. Compute the dot product $langle mathbf{g}_{cont}, mathbf{g}_{intent} rangle$. Where negative, execute vector subtraction to decouple opposing gradient forces.
- Parameter Update: Aggregate projected gradient tensors into the operational gradient field and execute the AdamW/Lion optimizer step.
Empirical Impact on Keyword Cluster Tool Accuracy
Integrating PCGrad and DWA directly improves the downstream topological representations generated for search engine queries. When deploying embeddings generated by this framework into downstream agglomerative hierarchical clustering (AHC) or HDBSCAN pipelines, metric performance shows marked divergence from standard baselines:
- Intent Purity Within Dense Semantic Neighborhoods: Unbalanced optimization frequently clusters commercial transactional terms (e.g., “enterprise keyword cluster tool pricing”) with broad conceptual queries (e.g., “how does keyword clustering work”) due to high cross-attention between token representations. The orthogonalized training prevents cross-entropy over-smoothing, maintaining clear hyper-plane separation between intents without fracturing the underlying topical similarity space.
- Enhanced Manifold Smoothness: By mitigating catastrophic interference between contrastive pulling/pushing forces and class boundary separation, the semantic space avoids local dimensionality collapse, yielding more uniform density distributions favorable for density-based clustering routines.
- Pareto Front Stabilization: Evaluation of training runs reveals that DWA dynamically allocates higher gradient budgets to $mathcal{L}_{cont}$ during early training phases—when negative pairs require aggressive separation—before re-allocating capacity to intent boundaries as the contrastive manifold stabilizes.
By enforcing microscopic geometric alignment via PCGrad and macroscopic adaptive capacity allocation via DWA, the multi-task bi-encoder produces high-fidelity embedding spaces engineered to satisfy both semantic depth and intent selectivity.