Judge-Aware Gated Multi-Task Learning

Updated 5 July 2026

The paper introduces a judge-aware gating mechanism that dynamically modulates shared and specific representations to enhance legal outcome predictions.
It employs learned judge embeddings and task-aware routing to reduce negative transfer and improve performance on ambiguous or infrequent classes.
The architecture integrates multi-task loss balancing and structural sparsity to offer both interpretability and efficiency over traditional generative models.

Judge-aware gated multi-task learning denotes a family of multi-task architectures in which a gating or routing mechanism modulates parameter sharing, specialization, or feature fusion using information that identifies or characterizes the evaluator, task rubric, or decision context. In the strictest sense represented in the current literature, the clearest implemented form is a legal outcome prediction architecture that conditions label-wise features on learned judge embeddings while jointly predicting coarse and fine-grained outcomes (Sójka et al., 25 Jun 2026). In a broader methodological sense, the surrounding design space includes task-aware hard gating and pruning in multi-task vision models, task-aware sparse mixture-of-experts routing, token-wise gated transfer in sequence tagging, and language-conditioned gating in multi-task reinforcement learning; conversely, recent MLLM-as-a-Judge work is explicitly judge-aware but does not implement gating in its base model (Bejnordi et al., 2024, Gupta et al., 2022, Ampomah et al., 2019, Arora, 7 Oct 2025, Wu et al., 12 Mar 2026).

1. Terminological scope and conceptual distinctions

The term judge-aware is not used uniformly across the cited literature. In legal outcome prediction, it refers literally to conditioning predictions on judge identity, with the stated goal of disentangling objective case facts from adjudicative context (Sójka et al., 25 Jun 2026). In multimodal evaluation models, judge-awareness refers instead to optimizing a model to behave as an evaluator rather than a general generator, enforced by task-specific judging prompts and a composite reward coupling structural adherence with correctness of the final verdict (Wu et al., 12 Mar 2026). In several neighboring strands of gated multi-task learning, the operative conditioning variable is not a judge at all, but rather a task identity, task embedding, auxiliary prediction stream, or natural-language task description (Bejnordi et al., 2024, Gupta et al., 2022, Ampomah et al., 2019, Arora, 7 Oct 2025).

This distinction matters because it separates three different design questions. The first is whether the system is evaluator-conditioned or identity-conditioned. The second is whether it actually uses gating or merely shared multi-task optimization. The third is whether the gating is dynamic at training only or remains input-dependent at inference. InterroGate, for example, learns per-task, per-layer, per-channel gates, but the paper explicitly states that it does not include any “judge-aware” mechanism and is better described as “task-aware” rather than “judge-aware” (Bejnordi et al., 2024). MT-RL-Judge is judge-aware in the evaluator sense, but “does not implement gating, routers, or task-conditioned adapters”; tasks are unified through prompts and a single shared policy optimized by GRPO (Wu et al., 12 Mar 2026).

System	Conditioning signal	Gating status
InterroGate	Task-conditioned gates $\alpha_t^\ell$	Implemented
MT-RL-Judge	Task-specific judging prompts and composite reward	Not implemented
UKET judge-aware model	Judge identity $j$ via learned embedding $e_j$	Implemented

A plausible implication is that the phrase judge-aware gated multi-task learning is best treated as a narrow architectural label in the legal setting and a broader comparative lens elsewhere. Under that lens, the shared problem is controlled conditional transfer: deciding when shared representations should dominate and when specialized pathways should override them.

A central architectural pattern is hard or soft selection between shared and specialized computation. InterroGate instantiates this most directly. For task $t$ and layer $\ell$ , it maintains shared features $s^\ell$ and task-specific features $\phi_t^\ell$ , and uses a binary per-channel selector

$\phi_t^{'\ell} = z_t^\ell \odot s^\ell + (1-z_t^\ell)\odot \phi_t^\ell,$

with $z_t^\ell = H(\sigma(\alpha_t^\ell)-0.5)$ trained by a straight-through estimator. Shared features are then rebuilt by a per-channel softmax mixing over tasks,

$s^{'\ell}[c] = \sum_{t=1}^{T}\pi_t^\ell[c]\phi_t^{'\ell}[c], \qquad \pi_t^\ell[c] = \frac{\exp(\beta^\ell[c,t])}{\sum_{u=1}^{T}\exp(\beta^\ell[c,u])},$

so the shared branch functions as a “communication bus” rather than a residual stream copied forward unchanged (Bejnordi et al., 2024).

Sparse MoE routing provides a second form. In "Sparsely Activated Mixture-of-Experts are Robust Multi-Task Learners" (Gupta et al., 2022), each Transformer FFN is replaced by a sparse MoE layer, and a task-aware gate with parameters $j$ 0 routes tokens to experts via

$j$ 1

With top-1 routing, only one expert is active per token, so FFN compute remains at parity with the dense baseline. The paper also explicitly maps this mechanism to a “judge-aware” view in which the gate acts as a judge scoring experts for a given input and task.

Sequence tagging introduces a third pattern: gated inter-task transfer over auxiliary predictions rather than over hidden-state experts. In the Gated Task Interaction network, auxiliary CRF outputs are embedded, contextualized, and then modulated token-wise by

$j$ 2

The gate is conditioned jointly on transformed auxiliary features and the main-task representation, so transfer is token-wise and context-sensitive rather than layer-wise and globally shared (Ampomah et al., 2019).

A fourth variant appears in multi-task RL. LEXPOL encodes natural-language task metadata with BERT, maps it through an MLP gate $j$ 3, and produces mixture weights

$j$ 4

which blend $j$ 5 sub-policies into a final action. In the base implementation the gate depends only on language metadata, not on a human or institutional judge, but the paper explicitly positions this gate as a natural locus for external evaluators or “judges” to modulate expert selection (Arora, 7 Oct 2025).

Taken together, these architectures differ in granularity—per-channel, per-token, per-expert, or per-policy—but share a common principle: multi-task learning is stabilized by making transfer conditional rather than mandatory.

3. Explicit judge-aware gated multi-task learning for legal outcome prediction

The most direct realization of judge-aware gated multi-task learning is the UK Employment Tribunal model introduced in "Towards Explainable Adjudicative Variance: Quantifying Judicial Discretion via Gated Multi-Task Learning" (Sójka et al., 25 Jun 2026). The problem is formalized with case text $j$ 6, judge identity $j$ 7, a fine-grained Detailed Case Outcome label $j$ 8, and a coarse General Case Outcome label $j$ 9 with four classes: Wins, Loses, Partly Wins, and Other. The DCO taxonomy contains 11 fine-grained classes with counts in 13,937 labeled UKET decisions: substantive win 3,463; substantive loss 3,341; default judgment 1,307; claimant partly wins 2,991; out of time 731; strike out non pursuit 579; strike out other 532; other procedural 398; lack of jurisdiction 377; withdrawal/settlement 193; other 698.

The backbone is Gemma-4 26B-A4B used in encoder mode, adapted by LoRA. For the hybrid B2 model, LoRA uses rank $e_j$ 0, scaling $e_j$ 1, and dropout 0.05, targeting all attention and MLP projections: $e_j$ 2 in attention and gate, up, down in the FFN. The architecture combines a shared encoder $e_j$ 3, label-wise attention (LWAN), a judge embedding, and two task heads. The shared encoder produces a pooled case representation $e_j$ 4. LWAN then constructs outcome-specific features $e_j$ 5 for each DCO label $e_j$ 6. Judge identity is embedded into a 48-dimensional vector $e_j$ 7, projected into the LWAN feature space as $e_j$ 8. Judge-aware modulation is implemented through

$e_j$ 9

followed by residual fusion

$t$ 0

The fused representation is then pooled into $t$ 1 and sent to two heads: a main GCO predictor and an auxiliary DCO predictor.

The training objective is a standard multi-task sum:

$t$ 2

with $t$ 3 reported as the best trade-off and an empirically stable range of $t$ 4; higher $t$ 5 degrades generalization. No explicit orthogonality penalty or gating entropy penalty is used in the reported configuration, so the total loss is simply

$t$ 6

Calibration is evaluated post hoc with Expected Calibration Error rather than trained directly.

The comparative results are framed around conditioning interface. The strongest generative SFT baseline, G4, injects both judge identity and DCO supervision through a single autoregressive channel and reaches 60.12 Macro-F1 / 73.26 Weighted-F1. The hybrid B2 model, which routes the same contextual signals through judge embeddings, gated fusion, and multi-task heads, reaches 65.21 / 75.55 and is reported as a new state of the art on this benchmark while training roughly an order of magnitude fewer parameters than the generative SFT baselines. Per-class GCO F1 on the test set ( $t$ 7) shows B2 at 81.40 for Wins, 81.80 for Loses, 56.60 for Partly Wins, and 41.10 for Other; the gains are concentrated on the ambiguous and rarest classes, especially Partly Wins and Other.

Model	Macro-F1	Weighted-F1
G4	60.12	73.26
B1	58.05	69.48
B2	65.21	75.55

The paper’s interpretability claims are tied directly to the gating mechanism. Judge-swap counterfactuals replace the original judge by $t$ 8 alternatives and compute the average KL divergence of GCO distributions; the mean KL increases from 0.055 in G4 to 0.141 in B2, indicating materially stronger and more discriminative conditioning on judge identity. Gate stratification further shows that technical disposals such as strike out non pursuit concentrate in the lowest gate quartile, whereas default judgment mass increases sharply in the highest quartile. The paper also reports that judge embeddings form semantically coherent clusters under t-SNE + k-means, with merit-focused and proceduralist regions. No formal significance tests or confidence intervals are reported.

4. Judge models, multimodal evaluation, and proposed judge-aware routing

The MLLM-as-a-Judge literature introduces a different notion of judge-awareness. MT-RL-Judge unifies six judging datasets—SeeTRUE, ImageReward, UnsafeBench, AGIN-Nat, AGIN-Tech, and AGIN-Rat—and optimizes a single multimodal judge based on Qwen3-VL-30B-A3B-Instruct using Group Relative Policy Optimization (Wu et al., 12 Mar 2026). Inputs are image(s) plus text, task identity is carried only through task-specific prompt templates, and the prompts require explicit reasoning traces in <thinking>...</thinking> followed by a boxed JSON label. The model’s judge-awareness is therefore enforced through prompting and reward design rather than through a routing mechanism.

Its reward is explicitly composite:

$t$ 9

and

$\ell$ 0

The unified RL objective is

$\ell$ 1

Policy optimization uses GRPO, with group-normalized advantages and KL regularization to the SFT reference policy. Tasks are mixed uniformly from the union of the six datasets; no task-specific weights or curriculum are reported.

The empirical profile is heterogeneous but instructive. On in-domain Macro-F1, MT-RL-Judge is best on SeeTRUE at 83.67 and second-best on several tasks: AGIN-Nat 81.63, AGIN-Rat 81.58, ImageReward 64.97. It is slightly below SFT-Single on UnsafeBench, where SFT-Single reaches 90.28 and MT-RL-Judge reaches 85.22. On OOD MJ-Bench pairwise evaluation, however, MT-RL-Judge outperforms both off-the-shelf prompting and unified SFT: Image–text Alignment 60.59 versus 59.41 and 55.82; Safety Judge 82.23 versus 73.07 and 49.40. The paper interprets this as evidence that unified multi-task RL produces transferable judging criteria.

Crucially, the paper states that “no mixture-of-experts or gating models are used in this work.” It then proposes, as an extension rather than a reported experiment, a judge-aware gating layer over lightweight experts such as LoRA adapters or small MLP heads. In that proposal, a gating network $\ell$ 2 produces mixture weights

$\ell$ 3

and expert outputs are combined as

$\ell$ 4

The proposed training objective augments the RL signal with gate regularizers:

$\ell$ 5

Because the judgment episodes are single-step, the return reduces to $\ell$ 6. This suggests a natural bridge between evaluator-centered RL and routed multi-task specialization, but in the current paper that bridge remains conceptual.

5. Interference mitigation, budget awareness, and the mechanics of specialization

A major motivation for gated multi-task learning is negative transfer. InterroGate addresses this structurally rather than through gradient surgery or loss reweighting. Its total objective is

$\ell$ 7

with

$\ell$ 8

and a hinge sparsity term

$\ell$ 9

The learned binary gate patterns are fixed at inference and used to prune task-specific and shared channels, collapsing the model into a static feed-forward architecture with no input-dependent routing. On NYUD-v2 with HRNet-18, InterroGate reaches $s^\ell$ 0 at 43.2 GFLOPs and 18.8M parameters, and $s^\ell$ 1 at 38.3 GFLOPs; the STL baseline is 65.1 GFLOPs with $s^\ell$ 2, while uniform/MTO methods are 24.5 GFLOPs with $s^\ell$ 3 to $s^\ell$ 4 (Bejnordi et al., 2024). On NYUD-v2 with ResNet-50, InterroGate reaches $s^\ell$ 5 at 916 GFLOPs and 95.4M parameters, outperforming Cross-stitch at $s^\ell$ 6 and approximately 1151 GFLOPs.

Task-aware sparse MoE reaches similar goals by restricting which weights receive gradients. On GLUE with a MiniLM encoder and $s^\ell$ 7 experts, the task-aware sparse MoE improves Small Tasks average to approximately 75.9 and All Tasks to approximately 82.5, versus dense MTL at approximately 74.9 and 82.1, and a shared-gate MoE at approximately 75.0 and 82.1. For unseen-task adaptation, the model reaches 90.0 on SciTail 1% and 90.3 on IMDB 1%, exceeding single-task, dense MTL, and shared-gate MoE baselines in those settings. When adding unrelated tasks, dense MTL reduces Small Tasks average by $s^\ell$ 8, whereas task-aware MoE improves it by $s^\ell$ 9 and raises All Tasks average from 87.32 to 87.83 (Gupta et al., 2022).

GTI provides a token-level analogue of the same phenomenon. On CoNLL-2003 NER with hidden size 200, single-task BiLSTM-CRF baselines yield mean F1 of 90.83–90.85, vanilla MTL reaches $\phi_t^\ell$ 0, pipeline MTL drops to $\phi_t^\ell$ 1, task interaction without the sigmoid gate reaches $\phi_t^\ell$ 2, and full GTI reaches $\phi_t^\ell$ 3. The paper interprets the gate’s contribution as suppressing harmful auxiliary signals and amplifying useful ones (Ampomah et al., 2019).

These results support a shared methodological claim: interference can be mitigated structurally by allowing tasks, tokens, or labels to reject unhelpful shared information. This is distinct from loss-balancing methods such as Uncertainty, DWA, Auto- $\phi_t^\ell$ 4, RLW, MGDA-UB, CAGrad, and PCGrad, which keep a shared encoder and modify optimization rather than the computation graph (Bejnordi et al., 2024).

6. Interpretability, limitations, and research directions

Interpretability is especially developed in the legal judge-aware model. The gate values and judge embeddings are used to localize where adjudicative context changes the prediction. Judge embeddings form clusters aligned with functional specializations, and gate magnitude stratification shows that accuracy increases monotonically with gate activation. Case studies identify examples in which the generative baseline collapses mixed or rare outcomes into LOSES, whereas the gated multi-task model recovers PARTLY WINS or OTHER by preserving alternative semantic pathways (Sójka et al., 25 Jun 2026). This is a stronger interpretability claim than in most task-aware gating work, where analysis typically focuses on routing statistics, capacity usage, or interference curves rather than institutionally meaningful latent factors.

Several limitations recur across the literature. InterroGate adds moderate training time and has only been demonstrated up to five tasks; its budget controls $\phi_t^\ell$ 5 and $\phi_t^\ell$ 6 steer but do not guarantee exact target FLOPs, and the current transformer instantiation is per-channel rather than structured at the head or token level (Bejnordi et al., 2024). MT-RL-Judge uses binary rewards and reports no statistical significance tests; the paper explicitly notes slight underperformance on UnsafeBench relative to SFT-Single and suggests that richer preference signals or safety-specific experts could help (Wu et al., 12 Mar 2026). Task-aware sparse MoE warns of gate under-training for small tasks, possible misrouting under domain shift, and scaling costs in expert storage even when FLOPs remain matched to a dense model (Gupta et al., 2022). LEXPOL depends on informative metadata and can face compute, memory, and credit-assignment difficulties as the number of experts grows; the hybrid LEXPOL + CARE model improves asymptotic performance but can underperform LEXPOL alone in low-sample regimes (Arora, 7 Oct 2025). GTI remains dependent on auxiliary-task quality and relatedness, and its auxiliary-to-main pathway is built from discrete one-best tags rather than a differentiable auxiliary posterior (Ampomah et al., 2019).

The open directions are correspondingly concrete. The legal paper proposes hierarchical MTL over GCO and DCO, per-class gates, mixture-of-experts routing with judge-aware priors, counterfactual judge-swap training, domain adaptation of judge embeddings, and calibration-aware training (Sójka et al., 25 Jun 2026). MT-RL-Judge proposes multimodal gates over experts or adapters conditioned on rubric cues and multimodal features (Wu et al., 12 Mar 2026). InterroGate points toward token-level gating in transformers (Bejnordi et al., 2024). Sparse MoE work suggests richer judge-aware embeddings, uncertainty-aware routing, and meta-learning for unseen tasks (Gupta et al., 2022). LEXPOL describes post-hoc multiplicative judge modulation, hierarchical gates, and auxiliary alignment losses to external evaluators (Arora, 7 Oct 2025).

A plausible synthesis is that judge-aware gated multi-task learning is becoming a general recipe for problems in which shared structure and context-specific discretion coexist. Where the context variable is a literal judge identity, as in legal prediction, gating provides a differentiable interface for controlled conditioning. Where the “judge” is an evaluator, rubric, or routing score, gating offers a mechanism for selective transfer and specialization. The common lesson across these settings is not merely that more context helps, but that the interface by which context enters the model—prompt tokens, shared parameters, residual fusion, or sparse routing—substantially determines both accuracy and interpretability.