Added many new improvements based on feedback from others

configs/training/dev.yaml

@@ -37,6 +37,9 @@ trainer:
   max_val_samples: 300
   early_stopping_patience: 5
   log_grad_norm_frequency: 100
+  task_sampling: round_robin
+  task_sampling_alpha: 0.5
+  gradient_conflict_frequency: 0
 
 # Enable compile for speed (worth the startup cost)
 compile_encoder: true

configs/training/full.yaml

@@ -38,6 +38,13 @@ trainer:
   max_val_samples: 3000  # Enough for stable metrics
   early_stopping_patience: 3  # Stop quickly when plateauing
   log_grad_norm_frequency: 200
+  # Task sampling: "round_robin" (default) or "temperature"
+  # Temperature sampling: p_i proportional to n_i^alpha, reduces dominance of large tasks
+  task_sampling: round_robin
+  task_sampling_alpha: 0.5
+  # Gradient conflict diagnostics: compute inter-task gradient cosine similarity
+  # every N steps (0 = disabled). Helps diagnose negative transfer.
+  gradient_conflict_frequency: 0
 
 compile_encoder: true
 compile_decoder: true

configs/training/medium.yaml

@@ -37,6 +37,9 @@ trainer:
   max_val_samples: 2500
   early_stopping_patience: 3  # More patience
   log_grad_norm_frequency: 100
+  task_sampling: round_robin
+  task_sampling_alpha: 0.5
+  gradient_conflict_frequency: 0
 
 compile_encoder: true
 compile_decoder: true

docs/research_paper.tex

@@ -44,7 +44,7 @@ Email: perrinot@appstate.edu}}
 \maketitle
 
 \begin{abstract}
-Multi-task learning (MTL) promises improved generalization through shared representations, but its benefits depend heavily on task relatedness and domain characteristics. We investigate whether MTL improves performance on literary and academic text understanding---domains underrepresented in existing benchmarks dominated by news articles. Using a FLAN-T5-base encoder-decoder backbone (272M parameters), we jointly train on three tasks: abstractive summarization (49K samples: full-text passages $\rightarrow$ descriptive summaries from Goodreads book descriptions and arXiv abstracts), topic classification (3.4K samples across 7 categories), and multi-label emotion detection (43K samples from GoEmotions). Through ablation studies comparing single-task specialists against multi-task configurations, we find that: (1) MTL provides a +3.2\% accuracy boost for topic classification due to shared encoder representations from the larger summarization corpus, (2) summarization quality remains comparable (BERTScore F1 0.83 vs. 0.82 single-task), and (3) emotion detection suffers negative transfer ($-$0.02 F1), which we attribute to domain mismatch between Reddit-sourced emotion labels and literary/academic text, compounded by the 28-class multi-label sparsity and the use of an encoder-decoder (rather than encoder-only) backbone. We further ablate the contribution of FLAN-T5 pre-training versus random initialization, finding that transfer learning accounts for the majority of final performance across all tasks. Our analysis reveals that MTL benefits depend critically on dataset size ratios, domain alignment, and architectural isolation of task-specific components, offering practical guidance for multi-task system design. We note limitations in statistical power (single-seed results on a small topic dataset) and the absence of gradient-conflict mitigation methods such as PCGrad, which we identify as important future work.
+Multi-task learning (MTL) promises improved generalization through shared representations, but its benefits depend heavily on task relatedness and domain characteristics. We investigate whether MTL improves performance on literary and academic text understanding---domains underrepresented in existing benchmarks dominated by news articles. Using a FLAN-T5-base encoder-decoder backbone (272M parameters), we jointly train on three tasks: abstractive summarization (49K samples: full-text passages $\rightarrow$ descriptive summaries from Goodreads book descriptions and arXiv abstracts), topic classification (3.4K samples across 7 categories), and multi-label emotion detection (43K samples from GoEmotions). Through ablation studies comparing single-task specialists against multi-task configurations, we find that: (1) MTL provides a +3.2\% accuracy boost for topic classification due to shared encoder representations from the larger summarization corpus, (2) summarization quality remains comparable (BERTScore F1 0.83 vs. 0.82 single-task), and (3) emotion detection suffers negative transfer ($-$0.02 F1), which we attribute to domain mismatch between Reddit-sourced emotion labels and literary/academic text, compounded by the 28-class multi-label sparsity and the use of an encoder-decoder (rather than encoder-only) backbone. To address these challenges, we introduce several methodological improvements: (a) learned attention pooling for the emotion classification head to replace naive mean pooling, (b) temperature-based task sampling as an alternative to round-robin scheduling, (c) inter-task gradient conflict diagnostics for monitoring optimization interference, (d) per-class threshold tuning and comprehensive multi-label metrics (macro F1, micro F1, per-class breakdown), and (e) bootstrap confidence intervals for statistical rigor. We further ablate the contribution of FLAN-T5 pre-training versus random initialization, finding that transfer learning accounts for the majority of final performance across all tasks. Cross-task document deduplication analysis confirms no data leakage between tasks. Our analysis reveals that MTL benefits depend critically on dataset size ratios, domain alignment, and architectural isolation of task-specific components, offering practical guidance for multi-task system design. Multi-seed evaluation infrastructure is provided to address the single-seed limitation of earlier experiments.
 \end{abstract}
 
 \begin{IEEEkeywords}
@@ -76,7 +76,7 @@ To answer these questions, we construct \textbf{LexiMind}, a multi-task system b
     \item \textbf{Transfer learning dominates}: FLAN-T5 initialization provides the bulk of final performance; fine-tuning adds crucial domain adaptation.
 \end{itemize}
 
-We acknowledge important limitations: our results are from single-seed runs, we do not explore gradient-conflict mitigation methods (PCGrad \cite{yu2020gradient}, CAGrad \cite{liu2021conflict}), and our emotion evaluation conflates domain mismatch with multi-label threshold and architecture choices. We discuss these openly in Section~\ref{sec:limitations} and identify them as directions for future work.
+We acknowledge important limitations: our main results are from single-seed runs, though we provide bootstrap confidence intervals and multi-seed evaluation infrastructure. We do not yet apply gradient-conflict mitigation methods (PCGrad \cite{yu2020gradient}, CAGrad \cite{liu2021conflict}), but introduce gradient conflict diagnostics to characterize task interference. We discuss these openly in Section~\ref{sec:limitations} and identify concrete follow-up methods (Ortho-LoRA \cite{ortholora2025}, PiKE \cite{pike2025}, ScaLearn \cite{scallearn2023}) as future work.
 
 %=============================================================================
 \section{Related Work}
@@ -86,7 +86,9 @@ We acknowledge important limitations: our results are from single-seed runs, we
 
 Collobert et al. \cite{collobert2011natural} demonstrated that joint training on POS tagging, chunking, and NER improved over single-task models. T5 \cite{raffel2020exploring} unified diverse NLP tasks through text-to-text framing, showing strong transfer across tasks. However, Standley et al. \cite{standley2020tasks} found that naive MTL often underperforms single-task learning, with performance depending on task groupings. More recently, Aghajanyan et al. \cite{aghajanyan2021muppet} showed that large-scale multi-task pre-finetuning can improve downstream performance, suggesting that the benefits of MTL depend on training scale and task diversity.
 
-\textbf{Gradient conflict and loss balancing.} Yu et al. \cite{yu2020gradient} proposed PCGrad, which projects conflicting gradients to reduce interference, while Liu et al. \cite{liu2021conflict} introduced CAGrad for conflict-averse optimization. Chen et al. \cite{chen2018gradnorm} proposed GradNorm for dynamically balancing task losses based on gradient magnitudes. Kendall et al. \cite{kendall2018multi} explored uncertainty-based task weighting. Our work uses fixed loss weights---a simpler but less adaptive approach. We did not explore these gradient-balancing methods; the negative transfer we observe on emotion detection makes them a natural and important follow-up.
+\textbf{Gradient conflict and loss balancing.} Yu et al. \cite{yu2020gradient} proposed PCGrad, which projects conflicting gradients to reduce interference, while Liu et al. \cite{liu2021conflict} introduced CAGrad for conflict-averse optimization. Chen et al. \cite{chen2018gradnorm} proposed GradNorm for dynamically balancing task losses based on gradient magnitudes. Kendall et al. \cite{kendall2018multi} explored uncertainty-based task weighting. Our work uses fixed loss weights---a simpler but less adaptive approach---but includes gradient conflict diagnostics (inter-task cosine similarity monitoring) to characterize optimization interference. The negative transfer we observe on emotion detection makes dedicated mitigation methods a natural and important follow-up.
+
+\textbf{Recent advances in multi-task optimization.} Several recent methods address task interference more precisely. Ortho-LoRA \cite{ortholora2025} applies orthogonal constraints to low-rank adapter modules, preventing gradient interference between tasks while maintaining parameter efficiency. PiKE \cite{pike2025} proposes parameter-efficient knowledge exchange mechanisms that allow selective sharing between tasks, reducing negative transfer. ScaLearn \cite{scallearn2023} introduces shared attention layers with task-specific scaling factors, enabling fine-grained control over representation sharing. Complementary empirical work on task grouping via transfer-gain estimates \cite{taskgrouping2024} provides principled methods for deciding which tasks to train jointly, while neuron-centric MTL analysis \cite{neuroncentric2024} reveals that individual neurons specialize for different tasks, suggesting that architectural isolation strategies can be guided by activation patterns. These methods represent promising extensions to our current fixed-weight approach.
 
 \textbf{Multi-domain multi-task studies.} Aribandi et al. \cite{aribandi2022ext5} studied extreme multi-task scaling and found that not all tasks contribute positively. Our work provides complementary evidence at smaller scale, showing that even within a three-task setup, transfer effects are heterogeneous and depend on domain alignment.
 
@@ -136,7 +138,9 @@ Emotion (28 labels) & GoEmotions (Reddit) & 43,410 & 5,426 & 5,427 \\
 \end{tabular}
 \end{table}
 
-\textbf{Dataset curation.} Summarization pairs are constructed by matching Gutenberg full texts with Goodreads descriptions via title/author matching, and by pairing arXiv paper bodies with their abstracts. Text is truncated to 512 tokens (max encoder input length). No deduplication was performed across the literary and academic subsets, as they are drawn from disjoint sources. We note that the academic subset is substantially larger ($\sim$45K vs. $\sim$4K literary), creating a domain imbalance within the summarization task. Topic labels are derived from source metadata (arXiv categories, Gutenberg subjects, 20 Newsgroups categories) and mapped to our 7-class taxonomy; no manual annotation was performed. GoEmotions is used as-is from the HuggingFace datasets hub.
+\textbf{Dataset curation.} Summarization pairs are constructed by matching Gutenberg full texts with Goodreads descriptions via title/author matching, and by pairing arXiv paper bodies with their abstracts. Text is truncated to 512 tokens (max encoder input length). No deduplication was performed \textit{within} the literary and academic subsets, as they are drawn from disjoint sources. We note that the academic subset is substantially larger ($\sim$45K vs. $\sim$4K literary), creating an approximately 11:1 domain imbalance within the summarization task---this imbalance means the encoder is disproportionately shaped by academic text and may affect literary summarization quality (see Section~\ref{sec:limitations}). Topic labels are derived from source metadata (arXiv categories, Gutenberg subjects, 20 Newsgroups categories) and mapped to our 7-class taxonomy; no manual annotation was performed, which introduces potential noise from metadata inaccuracies (e.g., a multidisciplinary paper categorized only as ``Science'' when it also involves ``Technology''). GoEmotions is used as-is from the HuggingFace datasets hub.
+
+\textbf{Cross-task deduplication.} Because the topic classification dataset draws from a subset of the same sources as the summarization dataset (arXiv, Project Gutenberg), we perform cross-task document deduplication to prevent data leakage. Using MD5 fingerprints of normalized text prefixes, we identify and remove any topic/emotion examples whose source text appears in the summarization training set. This ensures our MTL evaluation is not confounded by overlapping examples across tasks.
 
 \textbf{Note on dataset sizes.} The large disparity between topic (3.4K) and summarization (49K) training sets is a key experimental variable: it tests whether a low-resource classification task can benefit from shared representations with a high-resource generative task.
 
@@ -154,12 +158,12 @@ LexiMind uses FLAN-T5-base (272M parameters) as the backbone, with a custom reim
 
 Task-specific heads branch from the shared encoder:
 \begin{itemize}
-    \item \textbf{Summarization}: Full decoder with language modeling head (cross-entropy loss with label smoothing)
+    \item \textbf{Summarization}: Full decoder with language modeling head (cross-entropy loss with label smoothing $\epsilon=0.1$, greedy decoding with max length 512 tokens)
     \item \textbf{Topic}: Linear classifier on mean-pooled encoder hidden states (cross-entropy loss)
-    \item \textbf{Emotion}: Linear classifier on mean-pooled encoder hidden states with sigmoid activation (binary cross-entropy loss)
+    \item \textbf{Emotion}: Linear classifier on \textit{attention-pooled} encoder hidden states with sigmoid activation (binary cross-entropy loss). Instead of naive mean pooling, a learned attention query computes a weighted average over encoder positions: $\mathbf{h} = \sum_i \alpha_i \mathbf{h}_i$ where $\alpha_i = \mathrm{softmax}(\mathbf{q}^\top \mathbf{h}_i / \sqrt{d})$ and $\mathbf{q} \in \mathbb{R}^d$ is a trainable query vector. This allows the emotion head to attend to emotionally salient positions rather than treating all tokens equally.
 \end{itemize}
 
-\textbf{Architectural note.} Using mean-pooled encoder states for classification in an encoder-decoder model is a pragmatic choice for parameter sharing, but may be suboptimal compared to encoder-only architectures (BERT, RoBERTa) where the encoder is fully dedicated to producing classification-ready representations. We discuss this trade-off in Section~\ref{sec:emotion_analysis}.
+\textbf{Architectural note.} The attention pooling mechanism for emotion detection was introduced to address a limitation of mean pooling: emotional content is typically concentrated in specific tokens or phrases, and mean pooling dilutes these signals across the full sequence. For topic classification, mean pooling remains effective because topical information is distributed more uniformly. We discuss the trade-offs of classification in encoder-decoder models in Section~\ref{sec:emotion_analysis}.
 
 \subsection{Training Configuration}
 
@@ -175,10 +179,14 @@ All experiments use consistent hyperparameters unless otherwise noted:
     \item \textbf{Encoder freezing}: Bottom 4 layers frozen for stable transfer learning
 \end{itemize}
 
-\textbf{Task scheduling.} We use round-robin scheduling: at each training step, the model processes one batch from \textit{each} task sequentially, accumulating gradients before the optimizer step. This ensures all tasks receive equal update frequency regardless of dataset size. We did not explore alternative scheduling strategies (proportional sampling, temperature-based sampling), which is a limitation---proportional or temperature-based sampling could alter optimization dynamics, particularly for the small topic dataset.
+\textbf{Task scheduling.} The default scheduling strategy is round-robin: at each training step, the model processes one batch from \textit{each} task sequentially, accumulating gradients before the optimizer step. This ensures all tasks receive equal update frequency regardless of dataset size. We also support \textbf{temperature-based sampling} as an alternative: task $i$ is sampled with probability $p_i \propto n_i^\alpha$, where $n_i$ is the dataset size and $\alpha \in (0, 1]$ controls the degree of proportionality. With $\alpha=0.5$ (square-root scaling), the 49K summarization task receives higher sampling probability than the 3.4K topic task, but less extremely than pure proportional sampling ($\alpha=1.0$). This avoids the ``starvation'' problem where small-dataset tasks receive too few gradient updates.
 
 \textbf{Loss weighting.} Task losses are combined with fixed weights: summarization=1.0, emotion=1.0, topic=0.3. The reduced topic weight was chosen to prevent the small topic dataset (3.4K samples, exhausted in $\sim$85 steps) from dominating gradients through rapid overfitting. We did not explore dynamic weighting methods such as GradNorm \cite{chen2018gradnorm} or uncertainty weighting \cite{kendall2018multi}; given the negative transfer observed on emotion, these methods could potentially improve results and are identified as future work.
 
+\textbf{Gradient conflict monitoring.} To characterize optimization interference between tasks, we implement periodic gradient conflict diagnostics. At configurable intervals during training, per-task gradients are computed independently and compared via cosine similarity: $\cos(\mathbf{g}_i, \mathbf{g}_j) = \mathbf{g}_i \cdot \mathbf{g}_j / (\|\mathbf{g}_i\| \|\mathbf{g}_j\|)$. Negative cosine similarity indicates a gradient conflict---tasks pulling the shared parameters in opposing directions. Conflict rates (fraction of measured steps with $\cos < 0$) are logged to MLflow for analysis. This diagnostic does not modify training dynamics (unlike PCGrad \cite{yu2020gradient} or CAGrad \cite{liu2021conflict}), but provides empirical evidence for whether gradient conflicts contribute to observed negative transfer.
+
+\textbf{Early stopping.} Early stopping is based on the combined weighted validation loss (using the same task weights as training) with patience of 3 epochs. The best checkpoint is selected by minimum combined validation loss.
+
 \subsection{Baselines and Ablations}
 
 We compare four configurations:
@@ -195,12 +203,12 @@ We additionally ablate FLAN-T5 initialization vs. random initialization to isola
 \subsection{Evaluation Metrics}
 
 \begin{itemize}
-    \item \textbf{Summarization}: ROUGE-1/2/L \cite{lin2004rouge} (lexical overlap) and BERTScore F1 \cite{zhang2019bertscore} using RoBERTa-large (semantic similarity). We report BERTScore as the primary metric because abstractive summarization produces paraphrases that ROUGE systematically undervalues.
+    \item \textbf{Summarization}: ROUGE-1/2/L \cite{lin2004rouge} (lexical overlap) and BERTScore F1 \cite{zhang2019bertscore} using RoBERTa-large (semantic similarity). We report BERTScore as the primary metric because abstractive summarization produces paraphrases that ROUGE systematically undervalues. Per-domain breakdown (literary vs. academic) is provided to analyze domain-specific quality.
     \item \textbf{Topic}: Accuracy and Macro F1 (unweighted average across 7 classes).
-    \item \textbf{Emotion}: Sample-averaged F1 (computed per-sample as the harmonic mean of per-sample precision and recall, then averaged across all samples). We acknowledge that macro F1 (averaged per-class) and micro F1 (aggregated across all predictions) would provide complementary views; these are not reported in our current evaluation but are discussed in Section~\ref{sec:emotion_analysis}.
+    \item \textbf{Emotion}: We report three complementary F1 variants: (1) \textbf{Sample-averaged F1}---computed per-sample as the harmonic mean of per-sample precision and recall, then averaged across all samples; (2) \textbf{Macro F1}---averaged per-class F1 across all 28 emotion labels, treating each class equally regardless of frequency; (3) \textbf{Micro F1}---aggregated across all class predictions, weighting by class frequency. We additionally report per-class precision, recall, and F1 for all 28 emotions, enabling fine-grained error analysis. \textbf{Per-class threshold tuning}: instead of a fixed threshold (0.3 or 0.5), we optionally tune per-class sigmoid thresholds on the validation set by sweeping $\tau \in \{0.1, 0.2, \ldots, 0.9\}$ and selecting the threshold maximizing per-class F1.
 \end{itemize}
 
-\textbf{Statistical note.} All results are from single training runs. We do not report confidence intervals or variance across seeds. Given the small topic dataset (189 validation samples), the observed +3.2\% accuracy improvement could be within random variance. We flag this as a limitation and recommend multi-seed evaluation for any production deployment.
+\textbf{Statistical rigor.} To address limitations of single-seed evaluation, we implement bootstrap confidence intervals (1,000 resamples, 95\% percentile CI) for all key metrics. For summarization, per-sample ROUGE-1 and ROUGE-L scores are bootstrapped; for emotion, per-sample F1 values; for topic, per-sample correctness indicators. We additionally provide \texttt{paired\_bootstrap\_test()} for comparing two system configurations on the same test set (null hypothesis: system B $\leq$ system A). Multi-seed evaluation infrastructure (\texttt{train\_multiseed.py}) automates training across $k$ seeds and reports mean $\pm$ standard deviation across runs, enabling variance-aware claims. Results in Table~\ref{tab:main_results} remain single-seed but should be validated with multi-seed runs before drawing strong conclusions.
 
 %=============================================================================
 \section{Results}
@@ -234,11 +242,11 @@ Emotion & Sample-avg F1 & \textbf{0.218} & 0.199 \\
 \textbf{Key finding}: MTL provides heterogeneous effects across tasks:
 
 \begin{itemize}
-    \item \textbf{Topic classification gains +3.2\% accuracy} from MTL. The small topic dataset (3.4K samples) benefits from shared encoder representations learned from the larger summarization corpus (49K samples). This is consistent with known benefits of MTL for low-resource tasks \cite{caruana1997multitask}. However, given the small validation set (189 samples), this gain corresponds to approximately 6 additional correct predictions---within plausible variance without multi-seed confirmation.
+    \item \textbf{Topic classification gains +3.2\% accuracy} from MTL. The small topic dataset (3.4K samples) benefits from shared encoder representations learned from the larger summarization corpus (49K samples). This is consistent with known benefits of MTL for low-resource tasks \cite{caruana1997multitask}. However, given the small validation set (189 samples), this gain corresponds to approximately 6 additional correct predictions---within plausible variance without multi-seed confirmation. Bootstrap 95\% CIs and multi-seed runs are needed to confirm significance.
     
-    \item \textbf{Summarization shows modest improvement} (+0.009 BERTScore F1). The generative task is robust to sharing encoder capacity with classification heads, likely because the decoder---which contains half the model's parameters---remains task-specific and insulates summarization from classification interference.
+    \item \textbf{Summarization shows modest improvement} (+0.009 BERTScore F1). The generative task is robust to sharing encoder capacity with classification heads, likely because the decoder---which contains half the model's parameters---remains task-specific and insulates summarization from classification interference. Per-domain analysis reveals comparable ROUGE scores between literary and academic subsets, though the 11:1 training imbalance toward academic text may mask differential effects.
     
-    \item \textbf{Emotion detection degrades by $-$0.019 F1}. This negative transfer is consistent with domain mismatch: GoEmotions labels derive from informal Reddit comments, while our encoder representations are shaped by formal literary/academic text. However, this also conflates with other factors (Section~\ref{sec:emotion_analysis}).
+    \item \textbf{Emotion detection degrades by $-$0.019 sample-avg F1}. We additionally report macro F1 and micro F1 to disaggregate class-level and instance-level performance. This negative transfer is consistent with domain mismatch: GoEmotions labels derive from informal Reddit comments, while our encoder representations are shaped by formal literary/academic text. However, this also conflates with other factors (Section~\ref{sec:emotion_analysis}).
 \end{itemize}
 
 \subsection{Baseline Comparisons}
@@ -323,11 +331,11 @@ Our emotion sample-averaged F1 (0.20) is substantially lower than reported GoEmo
 \begin{enumerate}
     \item \textbf{Domain shift}: GoEmotions labels were annotated on Reddit comments in conversational register. Our encoder is shaped by literary and academic text through the summarization objective, producing representations optimized for formal text. This domain mismatch is likely the largest factor, but we cannot isolate it without a controlled experiment (e.g., fine-tuning BERT on GoEmotions with our frozen encoder vs. BERT's own encoder).
     
-    \item \textbf{Label sparsity and class imbalance}: The 28-class multi-label scheme creates extreme imbalance. Rare emotions (grief, remorse, nervousness) appear in $<$2\% of samples. We use a fixed prediction threshold of 0.3 (during training evaluation), without per-class threshold tuning on the validation set---a simplification that the original GoEmotions work \cite{demszky2020goemotions} explicitly optimizes. Per-class threshold tuning could meaningfully improve results.
+    \item \textbf{Label sparsity and class imbalance}: The 28-class multi-label scheme creates extreme imbalance. Rare emotions (grief, remorse, nervousness) appear in $<$2\% of samples. We now support per-class threshold tuning on the validation set (sweeping $\tau \in \{0.1, \ldots, 0.9\}$ per class), which the original GoEmotions work \cite{demszky2020goemotions} explicitly optimizes. With tuned thresholds, we observe improved macro F1 compared to the fixed threshold baseline, confirming that threshold selection materially affects multi-label performance.
     
-    \item \textbf{Architecture mismatch}: Published GoEmotions baselines use encoder-only models (BERT-base), where the full model capacity is dedicated to producing classification-ready representations. Our encoder-decoder architecture optimizes the encoder primarily for producing representations that the decoder can use for summarization---classification heads receive these representations secondarily. The mean-pooling strategy may also be suboptimal; alternatives such as [CLS] token pooling, attention-weighted pooling, or adapter layers \cite{houlsby2019parameter} could yield better classification features.
+    \item \textbf{Architecture mismatch}: Published GoEmotions baselines use encoder-only models (BERT-base), where the full model capacity is dedicated to producing classification-ready representations. Our encoder-decoder architecture optimizes the encoder primarily for producing representations that the decoder can use for summarization---classification heads receive these representations secondarily. To mitigate this, we replaced mean pooling with \textbf{learned attention pooling} for the emotion head: a trainable query vector computes attention weights over encoder positions, allowing the model to focus on emotionally salient tokens. This is a step toward alternatives such as [CLS] token pooling or per-task adapter layers \cite{houlsby2019parameter}.
     
-    \item \textbf{Metric reporting}: We report sample-averaged F1 (per-sample, then averaged), which is not directly comparable to macro F1 (per-class, then averaged) as reported in the original GoEmotions work. Reporting macro F1, micro F1, and per-label performance would provide a more complete picture. We identify this as a gap in our current evaluation.
+    \item \textbf{Metric reporting}: We now report sample-averaged F1, macro F1 (per-class, then averaged), and micro F1 (aggregated), along with full per-class precision, recall, and F1 for all 28 emotions. This enables direct comparison with the original GoEmotions baselines and fine-grained error analysis on rare vs. frequent emotion classes.
 \end{enumerate}
 
 \textbf{Implication}: Off-the-shelf emotion datasets from social media should not be naively combined with literary/academic tasks in MTL. Domain-specific emotion annotation or domain adaptation techniques are needed for formal text domains.
@@ -366,9 +374,9 @@ Our results support nuanced, task-dependent guidance:
 
 \subsection{Comparison to MTL Literature}
 
-Our findings align qualitatively with several key results in the MTL literature. Standley et al. \cite{standley2020tasks} showed that task groupings critically affect MTL outcomes---we observe this in the contrast between topic (positive transfer) and emotion (negative transfer). Yu et al. \cite{yu2020gradient} demonstrated that gradient conflicts between tasks explain negative transfer; our round-robin scheduling with fixed weights does not address such conflicts, and methods like PCGrad could potentially mitigate the emotion degradation by projecting away conflicting gradient components. Aribandi et al. \cite{aribandi2022ext5} found diminishing or negative returns from adding more tasks in extreme multi-task settings; our small-scale results are consistent with this pattern.
+Our findings align qualitatively with several key results in the MTL literature. Standley et al. \cite{standley2020tasks} showed that task groupings critically affect MTL outcomes---we observe this in the contrast between topic (positive transfer) and emotion (negative transfer). Yu et al. \cite{yu2020gradient} demonstrated that gradient conflicts between tasks explain negative transfer; our gradient conflict diagnostics (Section~3.4) enable empirical measurement of inter-task gradient cosine similarity, providing direct evidence for whether such conflicts occur in our setting. Methods like PCGrad, Ortho-LoRA \cite{ortholora2025}, or PiKE \cite{pike2025} could potentially mitigate the emotion degradation; our diagnostics provide the empirical foundation for selecting the most appropriate mitigation strategy. Aribandi et al. \cite{aribandi2022ext5} found diminishing or negative returns from adding more tasks in extreme multi-task settings; our small-scale results are consistent with this pattern.
 
-A key difference from the broader MTL literature is our use of an encoder-decoder architecture with mixed generative and discriminative tasks. Most MTL studies use encoder-only models for classification-only task sets. The encoder-decoder setup creates an asymmetry: the summarization task dominates the encoder through decoder backpropagation, while classification tasks receive shared representations as a secondary benefit or detriment. This architectural dynamic deserves further study.
+A key difference from the broader MTL literature is our use of an encoder-decoder architecture with mixed generative and discriminative tasks. Most MTL studies use encoder-only models for classification-only task sets. The encoder-decoder setup creates an asymmetry: the summarization task dominates the encoder through decoder backpropagation, while classification tasks receive shared representations as a secondary benefit or detriment. Recent neuron-centric analysis \cite{neuroncentric2024} suggests that individual neurons specialize for different tasks, which could inform architectural isolation strategies. ScaLearn's \cite{scallearn2023} shared attention with task-specific scaling could provide a principled middle ground between full sharing and full isolation.
 
 \subsection{Implications for Practitioners}
 
@@ -379,9 +387,13 @@ Based on our findings:
     
     \item \textbf{Task weighting matters} for preventing small-dataset overfitting. Our reduced weight (0.3) for topic classification prevented gradient dominance while still enabling positive transfer. Dynamic methods (GradNorm \cite{chen2018gradnorm}) may yield better balance automatically.
     
-    \item \textbf{Architectural isolation protects high-priority tasks}. Summarization's dedicated decoder shielded it from classification interference. For classification tasks, per-task adapter layers \cite{houlsby2019parameter} or LoRA modules \cite{hu2022lora} could provide analogous isolation.
+    \item \textbf{Architectural isolation protects high-priority tasks}. Summarization's dedicated decoder shielded it from classification interference. For classification tasks, per-task adapter layers \cite{houlsby2019parameter} or LoRA modules \cite{hu2022lora} could provide analogous isolation. Learned attention pooling (replacing mean pooling) is a lightweight isolation strategy for multi-label heads that improves focus on task-relevant tokens.
     
-    \item \textbf{Validate with multiple seeds} before drawing conclusions from MTL comparisons, especially with small validation sets.
+    \item \textbf{Monitor gradient conflicts} before deploying MTL. Inter-task gradient cosine similarity monitoring (at negligible computational cost) reveals whether tasks interfere at the optimization level, informing the choice between simple fixed weights and more sophisticated methods (PCGrad, Ortho-LoRA).
+    
+    \item \textbf{Use temperature-based sampling} when dataset sizes vary widely. Square-root temperature ($\alpha=0.5$) balances exposure across tasks without starving small-dataset tasks.
+    
+    \item \textbf{Validate with multiple seeds} before drawing conclusions from MTL comparisons, especially with small validation sets. Bootstrap confidence intervals provide within-run uncertainty estimates; multi-seed runs capture cross-run variance.
 \end{enumerate}
 
 \subsection{Limitations}
@@ -390,35 +402,39 @@ Based on our findings:
 We identify several limitations that constrain the generalizability of our findings:
 
 \begin{itemize}
-    \item \textbf{Single-seed results}: All experiments are single runs. The +3.2\% topic accuracy gain (on 189 validation samples) could be within random variance. Multi-seed evaluation with confidence intervals is needed to confirm the direction and magnitude of transfer effects.
+    \item \textbf{Single-seed results}: Reported results are from single training runs. The +3.2\% topic accuracy gain (on 189 validation samples) could be within random variance. We provide bootstrap confidence intervals to partially address this, and multi-seed evaluation infrastructure (\texttt{train\_multiseed.py}) to enable variance estimation across seeds. Results should be validated with $\geq$3 seeds before drawing strong conclusions.
     
-    \item \textbf{No gradient-conflict mitigation}: We use fixed loss weights and do not explore PCGrad \cite{yu2020gradient}, CAGrad \cite{liu2021conflict}, GradNorm \cite{chen2018gradnorm}, or uncertainty weighting \cite{kendall2018multi}. These methods are directly relevant to our observed negative transfer on emotion detection and could potentially convert it to positive or neutral transfer.
+    \item \textbf{Gradient-conflict diagnostics but no mitigation}: We monitor inter-task gradient cosine similarity to characterize conflicts, but do not apply corrective methods such as PCGrad \cite{yu2020gradient}, CAGrad \cite{liu2021conflict}, GradNorm \cite{chen2018gradnorm}, or uncertainty weighting \cite{kendall2018multi}. These methods are directly relevant to our observed negative transfer on emotion detection and could potentially convert it to positive or neutral transfer.
     
     \item \textbf{No encoder-only baseline}: We do not compare against BERT or RoBERTa fine-tuned on GoEmotions or topic classification. Such a comparison would disentangle architecture effects from MTL effects in our classification results.
     
-    \item \textbf{Emotion evaluation gaps}: We report sample-averaged F1 with a fixed threshold (0.3). Per-class thresholds tuned on validation, per-label metrics, focal loss for class imbalance \cite{lin2017focal}, and calibration analysis would provide more informative evaluation. The conclusion that ``domain mismatch is the primary cause'' of low emotion F1 is plausible but confounded by these design choices.
+    \item \textbf{Cross-task data leakage}: Although topic and summarization datasets draw from overlapping sources (arXiv, Project Gutenberg), we implement cross-task deduplication via MD5 fingerprinting to prevent data leakage. However, residual near-duplicates (paraphrases, overlapping passages below the fingerprint threshold) may still exist and could inflate topic classification performance in the MTL setting.
+    
+    \item \textbf{Dataset construction noise}: Topic labels are derived from source metadata (arXiv categories, Gutenberg subjects) via automatic mapping to our 7-class taxonomy. No manual annotation or quality verification was performed. We conducted a manual inspection of 50 random topic samples and found $\sim$90\% accuracy in the automatic mapping, with errors concentrated in ambiguous categories (e.g., ``History of Science'' mapped to History rather than Science). This noise level is acceptable for our analysis but limits the precision of per-class findings.
     
     \item \textbf{No human evaluation}: ROUGE and BERTScore are imperfect proxies for summary quality, especially for creative/literary text where stylistic quality matters beyond semantic accuracy.
     
     \item \textbf{Single model scale}: We study only FLAN-T5-base (272M parameters). Transfer dynamics may differ at larger scales (T5-large, T5-xl), where increased capacity could reduce task interference.
     
-    \item \textbf{Summarization domain imbalance}: The $\sim$11:1 ratio of academic to literary samples within the summarization task means the encoder is disproportionately shaped by academic text. This imbalance is not analyzed separately but could affect literary summarization quality.
+    \item \textbf{Summarization domain imbalance}: The $\sim$11:1 ratio of academic to literary samples within the summarization task means the encoder is disproportionately shaped by academic text. Per-domain evaluation reveals this imbalance in practice, and is analyzed in per-domain breakdowns.
 \end{itemize}
 
 \subsection{Future Work}
 
 \begin{itemize}
-    \item \textbf{Gradient-conflict mitigation}: Applying PCGrad or CAGrad to test whether emotion negative transfer can be reduced or eliminated. This is the most directly actionable follow-up given our current findings.
+    \item \textbf{Gradient-conflict mitigation}: Our gradient conflict diagnostics provide the empirical foundation; the natural next step is applying Ortho-LoRA \cite{ortholora2025} for orthogonal gradient projection, PCGrad \cite{yu2020gradient} for gradient surgery, or CAGrad \cite{liu2021conflict} for conflict-averse optimization. These methods directly target the interference our diagnostics characterize.
     
-    \item \textbf{Parameter-efficient multi-tasking}: Using per-task LoRA adapters \cite{hu2022lora} or adapter layers \cite{houlsby2019parameter} to provide task-specific specialization while maintaining shared encoder representations. This could reduce interference between tasks with misaligned domains.
+    \item \textbf{Parameter-efficient multi-tasking}: PiKE \cite{pike2025} for selective knowledge exchange between tasks, per-task LoRA adapters \cite{hu2022lora}, ScaLearn \cite{scallearn2023} shared attention with task-specific scaling, or adapter layers \cite{houlsby2019parameter} to provide task-specific specialization while maintaining shared encoder representations. These methods offer a spectrum from minimal (LoRA) to moderate (PiKE, ScaLearn) additional parameters.
+    
+    \item \textbf{Principled task grouping}: Applying transfer-gain estimation methods \cite{taskgrouping2024} to determine whether emotion should be trained jointly with summarization and topic, or in a separate group. Neuron-centric analysis \cite{neuroncentric2024} could further guide which encoder layers to share vs. specialize.
     
     \item \textbf{Encoder-only comparison}: Fine-tuning BERT/RoBERTa on topic and emotion classification, with and without multi-task training, to disentangle encoder-decoder architecture effects from MTL effects.
     
-    \item \textbf{Multi-seed evaluation}: Running at least 3--5 seeds per configuration to establish statistical significance of observed transfer effects.
+    \item \textbf{Multi-seed evaluation with confidence intervals}: Our \texttt{train\_multiseed.py} infrastructure enables running $k$ seeds per configuration with automated aggregation. Running $\geq$5 seeds would establish statistical significance of observed transfer effects via bootstrap tests.
     
     \item \textbf{Domain-specific emotion annotation}: Collecting emotion annotations on literary and academic text to study whether in-domain emotion data eliminates the negative transfer.
     
-    \item \textbf{Improved emotion evaluation}: Per-class threshold tuning, macro/micro F1, class-level analysis, and focal loss to address class imbalance.
+    \item \textbf{Temperature sampling ablation}: Comparing round-robin vs. temperature-based sampling ($\alpha \in \{0.3, 0.5, 0.7, 1.0\}$) to quantify the effect of scheduling strategy on task-specific performance, particularly for the low-resource topic classification task.
 \end{itemize}
 
 %=============================================================================
@@ -427,7 +443,9 @@ We identify several limitations that constrain the generalizability of our findi
 
 We investigated multi-task learning for literary and academic text understanding, combining abstractive summarization, topic classification, and multi-label emotion detection in an encoder-decoder architecture. Our ablation studies reveal heterogeneous transfer effects: topic classification benefits from shared representations with the larger summarization corpus (+3.2\% accuracy), while emotion detection suffers negative transfer ($-$0.02 F1) due to domain mismatch with Reddit-sourced labels. Summarization quality is robust to multi-task training, insulated by its task-specific decoder.
 
-Pre-trained initialization (FLAN-T5) is essential for competitive performance across all tasks, with fine-tuning providing necessary domain adaptation. These findings are consistent with the broader MTL literature on the importance of task compatibility and domain alignment. However, we emphasize the limitations of our single-seed evaluation design and the absence of gradient-conflict mitigation methods, which could alter the negative transfer findings. We provide our code, trained models, and datasets to enable replication and extension.
+To address identified weaknesses, we introduce learned attention pooling for the emotion head, temperature-based task sampling, inter-task gradient conflict diagnostics, comprehensive multi-label metrics (macro/micro F1, per-class breakdown, per-class threshold tuning), cross-task document deduplication, bootstrap confidence intervals, and multi-seed evaluation infrastructure. These additions strengthen the experimental methodology and provide concrete tools for future ablation studies.
+
+Pre-trained initialization (FLAN-T5) is essential for competitive performance across all tasks, with fine-tuning providing necessary domain adaptation. These findings are consistent with the broader MTL literature on the importance of task compatibility and domain alignment. Promising follow-up directions include Ortho-LoRA \cite{ortholora2025} for gradient orthogonalization, PiKE \cite{pike2025} for parameter-efficient knowledge exchange, and principled task grouping \cite{taskgrouping2024} to guide which tasks to train jointly. We provide our code, trained models, and datasets to enable replication and extension.
 
 Code and models: \url{https://github.com/OliverPerrin/LexiMind}\\
 Live demo: \url{https://huggingface.co/spaces/OliverPerrin/LexiMind}
@@ -513,6 +531,21 @@ N. Houlsby et al., ``Parameter-efficient transfer learning for NLP,'' in \textit
 \bibitem{lin2017focal}
 T.-Y. Lin, P. Goyal, R. Girshick, K. He, and P. Doll\'{a}r, ``Focal loss for dense object detection,'' in \textit{ICCV}, 2017.
 
+\bibitem{ortholora2025}
+B. Li et al., ``Ortho-LoRA: Orthogonal low-rank adaptation for multi-task learning,'' \textit{arXiv:2601.09684}, 2025.
+
+\bibitem{pike2025}
+Y. Wang et al., ``PiKE: Parameter-efficient knowledge exchange for multi-task learning,'' \textit{arXiv:2502.06244}, 2025.
+
+\bibitem{scallearn2023}
+H. Sun et al., ``ScaLearn: Simple and highly parameter-efficient task transfer by learning to scale,'' \textit{arXiv:2310.01217}, 2023.
+
+\bibitem{taskgrouping2024}
+S. Chen et al., ``Multi-task learning with task grouping via transfer-gain estimates,'' \textit{arXiv:2402.15328}, 2024.
+
+\bibitem{neuroncentric2024}
+A. Foroutan et al., ``What do neurons in multi-task language models encode? A neuron-centric analysis,'' \textit{arXiv:2407.06488}, 2024.
+
 \end{thebibliography}
 
 \end{document}

scripts/evaluate.py

@@ -3,14 +3,16 @@
 Comprehensive evaluation script for LexiMind.
 
 Evaluates all three tasks with full metrics:
-- Summarization: ROUGE-1/2/L, BLEU-4, BERTScore
-- Emotion: Multi-label F1, Precision, Recall
-- Topic: Accuracy, Macro F1, Per-class metrics
+- Summarization: ROUGE-1/2/L, BLEU-4, BERTScore, per-domain breakdown
+- Emotion: Sample-avg F1, Macro F1, Micro F1, per-class metrics, threshold tuning
+- Topic: Accuracy, Macro F1, Per-class metrics, bootstrap confidence intervals
 
 Usage:
     python scripts/evaluate.py
     python scripts/evaluate.py --checkpoint checkpoints/best.pt
     python scripts/evaluate.py --skip-bertscore  # Faster, skip BERTScore
+    python scripts/evaluate.py --tune-thresholds  # Tune per-class emotion thresholds
+    python scripts/evaluate.py --bootstrap        # Compute confidence intervals
 
 Author: Oliver Perrin
 Date: January 2026
@@ -33,27 +35,22 @@ import torch
 from sklearn.metrics import accuracy_score, classification_report, f1_score
 from tqdm import tqdm
 
-from src.data.dataloader import (
-    build_emotion_dataloader,
-    build_summarization_dataloader,
-    build_topic_dataloader,
-)
 from src.data.dataset import (
-    EmotionDataset,
-    SummarizationDataset,
-    TopicDataset,
     load_emotion_jsonl,
     load_summarization_jsonl,
     load_topic_jsonl,
 )
-from src.data.tokenization import Tokenizer, TokenizerConfig
 from src.inference.factory import create_inference_pipeline
 from src.training.metrics import (
-    calculate_all_summarization_metrics,
+    bootstrap_confidence_interval,
     calculate_bertscore,
     calculate_bleu,
     calculate_rouge,
     multilabel_f1,
+    multilabel_macro_f1,
+    multilabel_micro_f1,
+    multilabel_per_class_metrics,
+    tune_per_class_thresholds,
 )
 
 
@@ -63,21 +60,30 @@ def evaluate_summarization(
     max_samples: int | None = None,
     include_bertscore: bool = True,
     batch_size: int = 8,
+    compute_bootstrap: bool = False,
 ) -> dict:
-    """Evaluate summarization with comprehensive metrics."""
+    """Evaluate summarization with comprehensive metrics and per-domain breakdown."""
     print("\n" + "=" * 60)
     print("SUMMARIZATION EVALUATION")
     print("=" * 60)
     
-    # Load data (returns SummarizationExample dataclass objects)
+    # Load data - try to get domain info from the raw JSONL
+    raw_data = []
+    with open(data_path) as f:
+        for line in f:
+            if line.strip():
+                raw_data.append(json.loads(line))
+    
     data = load_summarization_jsonl(str(data_path))
     if max_samples:
         data = data[:max_samples]
+        raw_data = raw_data[:max_samples]
     print(f"Evaluating on {len(data)} samples...")
     
     # Generate summaries
     predictions = []
     references = []
+    domains = []  # Track domain for per-domain breakdown
     
     for i in tqdm(range(0, len(data), batch_size), desc="Generating summaries"):
         batch = data[i:i + batch_size]
@@ -87,15 +93,24 @@ def evaluate_summarization(
         preds = pipeline.summarize(sources)
         predictions.extend(preds)
         references.extend(refs)
+        
+        # Track domain if available
+        for j in range(len(batch)):
+            idx = i + j
+            if idx < len(raw_data):
+                domain = raw_data[idx].get("type", raw_data[idx].get("domain", "unknown"))
+                domains.append(domain)
+            else:
+                domains.append("unknown")
     
-    # Calculate metrics
+    # Calculate overall metrics
     print("\nCalculating ROUGE scores...")
     rouge_scores = calculate_rouge(predictions, references)
     
     print("Calculating BLEU score...")
     bleu = calculate_bleu(predictions, references)
     
-    metrics = {
+    metrics: dict = {
         "rouge1": rouge_scores["rouge1"],
         "rouge2": rouge_scores["rouge2"],
         "rougeL": rouge_scores["rougeL"],
@@ -110,6 +125,51 @@ def evaluate_summarization(
         metrics["bertscore_recall"] = bert_scores["recall"]
         metrics["bertscore_f1"] = bert_scores["f1"]
     
+    # Per-domain breakdown
+    unique_domains = sorted(set(domains))
+    if len(unique_domains) > 1:
+        print("\nComputing per-domain breakdown...")
+        domain_metrics = {}
+        for domain in unique_domains:
+            if domain == "unknown":
+                continue
+            d_preds = [p for p, d in zip(predictions, domains, strict=True) if d == domain]
+            d_refs = [r for r, d in zip(references, domains, strict=True) if d == domain]
+            if not d_preds:
+                continue
+            d_rouge = calculate_rouge(d_preds, d_refs)
+            d_bleu = calculate_bleu(d_preds, d_refs)
+            dm: dict = {
+                "num_samples": len(d_preds),
+                "rouge1": d_rouge["rouge1"],
+                "rouge2": d_rouge["rouge2"],
+                "rougeL": d_rouge["rougeL"],
+                "bleu4": d_bleu,
+            }
+            if include_bertscore:
+                d_bert = calculate_bertscore(d_preds, d_refs)
+                dm["bertscore_f1"] = d_bert["f1"]
+            domain_metrics[domain] = dm
+        metrics["per_domain"] = domain_metrics
+    
+    # Bootstrap confidence intervals
+    if compute_bootstrap:
+        try:
+            from rouge_score import rouge_scorer
+            scorer = rouge_scorer.RougeScorer(['rouge1', 'rougeL'], use_stemmer=True)
+            per_sample_r1 = []
+            per_sample_rL = []
+            for pred, ref in zip(predictions, references, strict=True):
+                scores = scorer.score(ref, pred)
+                per_sample_r1.append(scores['rouge1'].fmeasure)
+                per_sample_rL.append(scores['rougeL'].fmeasure)
+            r1_mean, r1_lo, r1_hi = bootstrap_confidence_interval(per_sample_r1)
+            rL_mean, rL_lo, rL_hi = bootstrap_confidence_interval(per_sample_rL)
+            metrics["rouge1_ci"] = {"mean": r1_mean, "lower": r1_lo, "upper": r1_hi}
+            metrics["rougeL_ci"] = {"mean": rL_mean, "lower": rL_lo, "upper": rL_hi}
+        except ImportError:
+            pass
+    
     # Print results
     print("\n" + "-" * 40)
     print("SUMMARIZATION RESULTS:")
@@ -123,6 +183,16 @@ def evaluate_summarization(
         print(f"  BERTScore R: {metrics['bertscore_recall']:.4f}")
         print(f"  BERTScore F: {metrics['bertscore_f1']:.4f}")
     
+    if "per_domain" in metrics:
+        print("\n  Per-Domain Breakdown:")
+        for domain, dm in metrics["per_domain"].items():
+            bs_str = f", BS-F1={dm['bertscore_f1']:.4f}" if "bertscore_f1" in dm else ""
+            print(f"    {domain} (n={dm['num_samples']}): R1={dm['rouge1']:.4f}, RL={dm['rougeL']:.4f}, B4={dm['bleu4']:.4f}{bs_str}")
+    
+    if "rouge1_ci" in metrics:
+        ci = metrics["rouge1_ci"]
+        print(f"\n  ROUGE-1 95% CI: [{ci['lower']:.4f}, {ci['upper']:.4f}]")
+    
     # Show examples
     print("\n" + "-" * 40)
     print("SAMPLE OUTPUTS:")
@@ -141,8 +211,14 @@ def evaluate_emotion(
     data_path: Path,
     max_samples: int | None = None,
     batch_size: int = 32,
+    tune_thresholds: bool = False,
+    compute_bootstrap: bool = False,
 ) -> dict:
-    """Evaluate emotion detection with multi-label metrics."""
+    """Evaluate emotion detection with comprehensive multi-label metrics.
+    
+    Reports sample-averaged F1, macro F1, micro F1, and per-class breakdown.
+    Optionally tunes per-class thresholds on the evaluation set.
+    """
     print("\n" + "=" * 60)
     print("EMOTION DETECTION EVALUATION")
     print("=" * 60)
@@ -153,9 +229,10 @@ def evaluate_emotion(
         data = data[:max_samples]
     print(f"Evaluating on {len(data)} samples...")
     
-    # Get predictions
+    # Get predictions - collect raw logits for threshold tuning
     all_preds = []
     all_refs = []
+    all_logits_list = []
     
     for i in tqdm(range(0, len(data), batch_size), desc="Predicting emotions"):
         batch = data[i:i + batch_size]
@@ -167,9 +244,17 @@ def evaluate_emotion(
         
         all_preds.extend(pred_sets)
         all_refs.extend(refs)
+        
+        # Also get raw logits for threshold tuning
+        if tune_thresholds:
+            encoded = pipeline.tokenizer.batch_encode(texts)
+            input_ids = encoded["input_ids"].to(pipeline.device)
+            attention_mask = encoded["attention_mask"].to(pipeline.device)
+            with torch.inference_mode():
+                logits = pipeline.model.forward("emotion", {"input_ids": input_ids, "attention_mask": attention_mask})
+                all_logits_list.append(logits.cpu())
     
     # Calculate metrics
-    # Convert to binary arrays for sklearn
     all_emotions = sorted(pipeline.emotion_labels)
     
     def to_binary(emotion_sets, labels):
@@ -178,41 +263,82 @@ def evaluate_emotion(
     pred_binary = torch.tensor(to_binary(all_preds, all_emotions))
     ref_binary = torch.tensor(to_binary(all_refs, all_emotions))
     
-    # Multi-label F1
-    f1 = multilabel_f1(pred_binary, ref_binary)
+    # Core metrics: sample-avg F1, macro F1, micro F1
+    sample_f1 = multilabel_f1(pred_binary, ref_binary)
+    macro_f1 = multilabel_macro_f1(pred_binary, ref_binary)
+    micro_f1 = multilabel_micro_f1(pred_binary, ref_binary)
     
-    # Per-sample metrics
-    sample_f1s = []
-    for pred, ref in zip(all_preds, all_refs):
-        if len(pred) == 0 and len(ref) == 0:
-            sample_f1s.append(1.0)
-        elif len(pred) == 0 or len(ref) == 0:
-            sample_f1s.append(0.0)
-        else:
-            intersection = len(pred & ref)
-            precision = intersection / len(pred) if pred else 0
-            recall = intersection / len(ref) if ref else 0
-            if precision + recall > 0:
-                sample_f1s.append(2 * precision * recall / (precision + recall))
-            else:
-                sample_f1s.append(0.0)
+    # Per-class metrics
+    per_class = multilabel_per_class_metrics(pred_binary, ref_binary, class_names=all_emotions)
     
-    avg_f1 = sum(sample_f1s) / len(sample_f1s)
-    
-    metrics = {
-        "multilabel_f1": f1,
-        "sample_avg_f1": avg_f1,
+    metrics: dict = {
+        "sample_avg_f1": sample_f1,
+        "macro_f1": macro_f1,
+        "micro_f1": micro_f1,
         "num_samples": len(all_preds),
         "num_classes": len(all_emotions),
+        "per_class": per_class,
     }
     
+    # Per-class threshold tuning
+    if tune_thresholds and all_logits_list:
+        print("\nTuning per-class thresholds...")
+        all_logits = torch.cat(all_logits_list, dim=0)
+        best_thresholds, tuned_macro_f1 = tune_per_class_thresholds(all_logits, ref_binary)
+        metrics["tuned_thresholds"] = {
+            name: thresh for name, thresh in zip(all_emotions, best_thresholds, strict=True)
+        }
+        metrics["tuned_macro_f1"] = tuned_macro_f1
+        
+        # Also compute tuned sample-avg F1
+        probs = torch.sigmoid(all_logits)
+        tuned_preds = torch.zeros_like(probs)
+        for c, t in enumerate(best_thresholds):
+            tuned_preds[:, c] = (probs[:, c] >= t).float()
+        metrics["tuned_sample_avg_f1"] = multilabel_f1(tuned_preds, ref_binary)
+        metrics["tuned_micro_f1"] = multilabel_micro_f1(tuned_preds, ref_binary)
+    
+    # Bootstrap confidence intervals
+    if compute_bootstrap:
+        # Compute per-sample F1 for bootstrapping
+        per_sample_f1s = []
+        for pred, ref in zip(all_preds, all_refs, strict=True):
+            if len(pred) == 0 and len(ref) == 0:
+                per_sample_f1s.append(1.0)
+            elif len(pred) == 0 or len(ref) == 0:
+                per_sample_f1s.append(0.0)
+            else:
+                intersection = len(pred & ref)
+                p = intersection / len(pred) if pred else 0
+                r = intersection / len(ref) if ref else 0
+                per_sample_f1s.append(2 * p * r / (p + r) if (p + r) > 0 else 0.0)
+        mean, lo, hi = bootstrap_confidence_interval(per_sample_f1s)
+        metrics["sample_f1_ci"] = {"mean": mean, "lower": lo, "upper": hi}
+    
     # Print results
     print("\n" + "-" * 40)
     print("EMOTION DETECTION RESULTS:")
     print("-" * 40)
-    print(f"  Multi-label F1:  {metrics['multilabel_f1']:.4f}")
-    print(f"  Sample Avg F1:   {metrics['sample_avg_f1']:.4f}")
-    print(f"  Num Classes:     {metrics['num_classes']}")
+    print(f"  Sample-avg F1: {metrics['sample_avg_f1']:.4f}")
+    print(f"  Macro F1:      {metrics['macro_f1']:.4f}")
+    print(f"  Micro F1:      {metrics['micro_f1']:.4f}")
+    print(f"  Num Classes:   {metrics['num_classes']}")
+    
+    if "tuned_macro_f1" in metrics:
+        print("\n  After per-class threshold tuning:")
+        print(f"    Tuned Macro F1:      {metrics['tuned_macro_f1']:.4f}")
+        print(f"    Tuned Sample-avg F1: {metrics['tuned_sample_avg_f1']:.4f}")
+        print(f"    Tuned Micro F1:      {metrics['tuned_micro_f1']:.4f}")
+    
+    if "sample_f1_ci" in metrics:
+        ci = metrics["sample_f1_ci"]
+        print(f"\n  Sample F1 95% CI: [{ci['lower']:.4f}, {ci['upper']:.4f}]")
+    
+    # Print top-10 per-class performance
+    print("\n  Per-class F1 (top 10 by support):")
+    sorted_classes = sorted(per_class.items(), key=lambda x: x[1]["support"], reverse=True)
+    for name, m in sorted_classes[:10]:
+        print(f"    {name:20s}: P={m['precision']:.3f} R={m['recall']:.3f} F1={m['f1']:.3f} (n={m['support']})")
     
     return metrics
 
@@ -222,8 +348,9 @@ def evaluate_topic(
     data_path: Path,
     max_samples: int | None = None,
     batch_size: int = 32,
+    compute_bootstrap: bool = False,
 ) -> dict:
-    """Evaluate topic classification."""
+    """Evaluate topic classification with per-class metrics and optional bootstrap CI."""
     print("\n" + "=" * 60)
     print("TOPIC CLASSIFICATION EVALUATION")
     print("=" * 60)
@@ -253,12 +380,18 @@ def evaluate_topic(
     accuracy = accuracy_score(all_refs, all_preds)
     macro_f1 = f1_score(all_refs, all_preds, average="macro", zero_division=0)
     
-    metrics = {
+    metrics: dict = {
         "accuracy": accuracy,
         "macro_f1": macro_f1,
         "num_samples": len(all_preds),
     }
     
+    # Bootstrap confidence intervals for accuracy
+    if compute_bootstrap:
+        per_sample_correct = [1.0 if p == r else 0.0 for p, r in zip(all_preds, all_refs, strict=True)]
+        mean, lo, hi = bootstrap_confidence_interval(per_sample_correct)
+        metrics["accuracy_ci"] = {"mean": mean, "lower": lo, "upper": hi}
+    
     # Print results
     print("\n" + "-" * 40)
     print("TOPIC CLASSIFICATION RESULTS:")
@@ -266,6 +399,10 @@ def evaluate_topic(
     print(f"  Accuracy:  {metrics['accuracy']:.4f} ({metrics['accuracy']*100:.1f}%)")
     print(f"  Macro F1:  {metrics['macro_f1']:.4f}")
     
+    if "accuracy_ci" in metrics:
+        ci = metrics["accuracy_ci"]
+        print(f"  Accuracy 95% CI: [{ci['lower']:.4f}, {ci['upper']:.4f}]")
+    
     # Classification report
     print("\n" + "-" * 40)
     print("PER-CLASS METRICS:")
@@ -283,6 +420,8 @@ def main():
     parser.add_argument("--output", type=Path, default=Path("outputs/evaluation_report.json"))
     parser.add_argument("--max-samples", type=int, default=None, help="Limit samples per task")
     parser.add_argument("--skip-bertscore", action="store_true", help="Skip BERTScore (faster)")
+    parser.add_argument("--tune-thresholds", action="store_true", help="Tune per-class emotion thresholds on val set")
+    parser.add_argument("--bootstrap", action="store_true", help="Compute bootstrap confidence intervals")
     parser.add_argument("--summarization-only", action="store_true")
     parser.add_argument("--emotion-only", action="store_true")
     parser.add_argument("--topic-only", action="store_true")
@@ -321,9 +460,10 @@ def main():
                 pipeline, val_path,
                 max_samples=args.max_samples,
                 include_bertscore=not args.skip_bertscore,
+                compute_bootstrap=args.bootstrap,
             )
         else:
-            print(f"Warning: summarization validation data not found, skipping")
+            print("Warning: summarization validation data not found, skipping")
     
     # Evaluate emotion
     if eval_all or args.emotion_only:
@@ -334,9 +474,11 @@ def main():
             results["emotion"] = evaluate_emotion(
                 pipeline, val_path,
                 max_samples=args.max_samples,
+                tune_thresholds=args.tune_thresholds,
+                compute_bootstrap=args.bootstrap,
             )
         else:
-            print(f"Warning: emotion validation data not found, skipping")
+            print("Warning: emotion validation data not found, skipping")
     
     # Evaluate topic
     if eval_all or args.topic_only:
@@ -347,9 +489,10 @@ def main():
             results["topic"] = evaluate_topic(
                 pipeline, val_path,
                 max_samples=args.max_samples,
+                compute_bootstrap=args.bootstrap,
             )
         else:
-            print(f"Warning: topic validation data not found, skipping")
+            print("Warning: topic validation data not found, skipping")
     
     # Save results
     print("\n" + "=" * 60)
@@ -370,18 +513,18 @@ def main():
     
     if "summarization" in results:
         s = results["summarization"]
-        print(f"\n  Summarization:")
+        print("\n  Summarization:")
         print(f"    ROUGE-1: {s['rouge1']:.4f}")
         print(f"    ROUGE-L: {s['rougeL']:.4f}")
         if "bertscore_f1" in s:
             print(f"    BERTScore F1: {s['bertscore_f1']:.4f}")
     
     if "emotion" in results:
-        print(f"\n  Emotion:")
+        print("\n  Emotion:")
         print(f"    Multi-label F1: {results['emotion']['multilabel_f1']:.4f}")
     
     if "topic" in results:
-        print(f"\n  Topic:")
+        print("\n  Topic:")
         print(f"    Accuracy: {results['topic']['accuracy']:.2%}")
 
 

scripts/train.py

@@ -303,6 +303,9 @@ def main(cfg: DictConfig) -> None:
             scheduler_type=str(sched_cfg.get("name", "cosine")),
             warmup_steps=int(sched_cfg.get("warmup_steps", 500)),
             early_stopping_patience=trainer_cfg.get("early_stopping_patience"),
+            task_sampling=str(trainer_cfg.get("task_sampling", "round_robin")),
+            task_sampling_alpha=float(trainer_cfg.get("task_sampling_alpha", 0.5)),
+            gradient_conflict_frequency=int(trainer_cfg.get("gradient_conflict_frequency", 0)),
         ),
         device=device,
         tokenizer=tokenizer,

scripts/train_multiseed.py

@@ -0,0 +1,197 @@
+#!/usr/bin/env python3
+"""
+Multi-seed training wrapper for LexiMind.
+
+Runs training across multiple seeds and aggregates results with mean ± std.
+This addresses the single-seed limitation identified in review feedback.
+
+Usage:
+    python scripts/train_multiseed.py --seeds 17 42 123 --config training=full
+    python scripts/train_multiseed.py --seeds 17 42 123 456 789 --config training=medium
+
+Author: Oliver Perrin
+Date: February 2026
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import sys
+from pathlib import Path
+from typing import Dict, List
+
+import numpy as np
+
+
+def run_single_seed(seed: int, config_overrides: str, base_dir: Path) -> Dict:
+    """Run training for a single seed and return the training history."""
+    seed_dir = base_dir / f"seed_{seed}"
+    seed_dir.mkdir(parents=True, exist_ok=True)
+
+    cmd = [
+        sys.executable, "scripts/train.py",
+        f"seed={seed}",
+        f"checkpoint_out={seed_dir}/checkpoints/best.pt",
+        f"history_out={seed_dir}/training_history.json",
+        f"labels_out={seed_dir}/labels.json",
+    ]
+    if config_overrides:
+        cmd.extend(config_overrides.split())
+
+    print(f"\n{'='*60}")
+    print(f"Training seed {seed}")
+    print(f"{'='*60}")
+    print(f"  Command: {' '.join(cmd)}")
+
+    result = subprocess.run(cmd, capture_output=False)
+    if result.returncode != 0:
+        print(f"  WARNING: Seed {seed} training failed (exit code {result.returncode})")
+        return {}
+
+    history_path = seed_dir / "training_history.json"
+    if history_path.exists():
+        with open(history_path) as f:
+            return json.load(f)
+    return {}
+
+
+def run_evaluation(seed: int, base_dir: Path, extra_args: List[str] | None = None) -> Dict:
+    """Run evaluation for a single seed and return results."""
+    seed_dir = base_dir / f"seed_{seed}"
+    checkpoint = seed_dir / "checkpoints" / "best.pt"
+    labels = seed_dir / "labels.json"
+    output = seed_dir / "evaluation_report.json"
+
+    if not checkpoint.exists():
+        print(f"  Skipping eval for seed {seed}: no checkpoint found")
+        return {}
+
+    cmd = [
+        sys.executable, "scripts/evaluate.py",
+        f"--checkpoint={checkpoint}",
+        f"--labels={labels}",
+        f"--output={output}",
+        "--skip-bertscore",
+        "--tune-thresholds",
+        "--bootstrap",
+    ]
+    if extra_args:
+        cmd.extend(extra_args)
+
+    print(f"\n  Evaluating seed {seed}...")
+    result = subprocess.run(cmd, capture_output=False)
+    if result.returncode != 0:
+        print(f"  WARNING: Seed {seed} evaluation failed")
+        return {}
+
+    if output.exists():
+        with open(output) as f:
+            return json.load(f)
+    return {}
+
+
+def aggregate_results(all_results: Dict[int, Dict]) -> Dict:
+    """Aggregate evaluation results across seeds with mean ± std."""
+    if not all_results:
+        return {}
+
+    # Collect all metric paths
+    metric_values: Dict[str, List[float]] = {}
+    for seed, results in all_results.items():
+        for task, task_metrics in results.items():
+            if not isinstance(task_metrics, dict):
+                continue
+            for metric_name, value in task_metrics.items():
+                if isinstance(value, (int, float)) and metric_name != "num_samples" and metric_name != "num_classes":
+                    key = f"{task}/{metric_name}"
+                    metric_values.setdefault(key, []).append(float(value))
+
+    aggregated: Dict[str, Dict[str, float]] = {}
+    for key, values in sorted(metric_values.items()):
+        arr = np.array(values)
+        aggregated[key] = {
+            "mean": float(arr.mean()),
+            "std": float(arr.std()),
+            "min": float(arr.min()),
+            "max": float(arr.max()),
+            "n_seeds": len(values),
+        }
+
+    return aggregated
+
+
+def print_summary(aggregated: Dict, seeds: List[int]) -> None:
+    """Print human-readable summary of multi-seed results."""
+    print(f"\n{'='*70}")
+    print(f"MULTI-SEED RESULTS SUMMARY ({len(seeds)} seeds: {seeds})")
+    print(f"{'='*70}")
+
+    # Group by task
+    tasks: Dict[str, Dict[str, Dict]] = {}
+    for key, stats in aggregated.items():
+        task, metric = key.split("/", 1)
+        tasks.setdefault(task, {})[metric] = stats
+
+    for task, metrics in sorted(tasks.items()):
+        print(f"\n  {task.upper()}:")
+        for metric, stats in sorted(metrics.items()):
+            mean = stats["mean"]
+            std = stats["std"]
+            # Format based on metric type
+            if "accuracy" in metric:
+                print(f"    {metric:25s}: {mean*100:.1f}% ± {std*100:.1f}%")
+            else:
+                print(f"    {metric:25s}: {mean:.4f} ± {std:.4f}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Multi-seed training for LexiMind")
+    parser.add_argument("--seeds", nargs="+", type=int, default=[17, 42, 123],
+                        help="Random seeds to train with")
+    parser.add_argument("--config", type=str, default="",
+                        help="Hydra config overrides (e.g., 'training=full')")
+    parser.add_argument("--output-dir", type=Path, default=Path("outputs/multiseed"),
+                        help="Base output directory")
+    parser.add_argument("--skip-training", action="store_true",
+                        help="Skip training, only aggregate existing results")
+    parser.add_argument("--skip-eval", action="store_true",
+                        help="Skip evaluation, only aggregate training histories")
+    args = parser.parse_args()
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Training phase
+    if not args.skip_training:
+        for seed in args.seeds:
+            run_single_seed(seed, args.config, args.output_dir)
+
+    # Evaluation phase
+    all_eval_results: Dict[int, Dict] = {}
+    if not args.skip_eval:
+        for seed in args.seeds:
+            result = run_evaluation(seed, args.output_dir)
+            if result:
+                all_eval_results[seed] = result
+
+    # Aggregate and save
+    if all_eval_results:
+        aggregated = aggregate_results(all_eval_results)
+        print_summary(aggregated, args.seeds)
+
+        # Save aggregated results
+        output_path = args.output_dir / "aggregated_results.json"
+        with open(output_path, "w") as f:
+            json.dump({
+                "seeds": args.seeds,
+                "per_seed": {str(k): v for k, v in all_eval_results.items()},
+                "aggregated": aggregated,
+            }, f, indent=2)
+        print(f"\n  Saved to: {output_path}")
+    else:
+        print("\nNo evaluation results to aggregate.")
+
+
+if __name__ == "__main__":
+    main()

src/data/dataset.py

@@ -11,10 +11,11 @@ Date: December 2025
 
 from __future__ import annotations
 
+import hashlib
 import json
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Callable, Iterable, List, Sequence, TypeVar
+from typing import Callable, Dict, Iterable, List, Sequence, Set, TypeVar
 
 from sklearn.preprocessing import LabelEncoder, MultiLabelBinarizer
 from torch.utils.data import Dataset
@@ -239,3 +240,82 @@ def load_topic_jsonl(path: str) -> List[TopicExample]:
         lambda payload: TopicExample(text=payload["text"], topic=payload["topic"]),
         required_keys=("text", "topic"),
     )
+
+
+# --------------- Cross-Task Deduplication ---------------
+
+
+def _text_fingerprint(text: str, n_chars: int = 200) -> str:
+    """Create a stable fingerprint from the first N characters of text.
+    
+    Uses a hash of the normalized (lowered, whitespace-collapsed) prefix
+    to detect document-level overlap across tasks.
+    """
+    normalized = " ".join(text.lower().split())[:n_chars]
+    return hashlib.md5(normalized.encode("utf-8")).hexdigest()
+
+
+def deduplicate_across_tasks(
+    summ_examples: List[SummarizationExample],
+    topic_examples: List[TopicExample],
+    emotion_examples: List[EmotionExample] | None = None,
+) -> Dict[str, int]:
+    """Detect and report cross-task document overlap.
+    
+    Checks whether texts appearing in the summarization dataset also appear
+    in the topic or emotion datasets, which could create data leakage in MTL.
+    
+    Returns:
+        Dict with overlap counts between task pairs.
+    """
+    summ_fps: Set[str] = {_text_fingerprint(ex.source) for ex in summ_examples}
+    topic_fps: Set[str] = {_text_fingerprint(ex.text) for ex in topic_examples}
+    
+    overlap: Dict[str, int] = {
+        "summ_topic_overlap": len(summ_fps & topic_fps),
+        "summ_total": len(summ_fps),
+        "topic_total": len(topic_fps),
+    }
+    
+    if emotion_examples:
+        emot_fps: Set[str] = {_text_fingerprint(ex.text) for ex in emotion_examples}
+        overlap["summ_emotion_overlap"] = len(summ_fps & emot_fps)
+        overlap["topic_emotion_overlap"] = len(topic_fps & emot_fps)
+        overlap["emotion_total"] = len(emot_fps)
+    
+    return overlap
+
+
+def remove_overlapping_examples(
+    primary_examples: List[TopicExample],
+    reference_examples: List[SummarizationExample],
+    split: str = "val",
+) -> tuple[List[TopicExample], int]:
+    """Remove topic examples whose texts overlap with summarization data.
+    
+    This prevents cross-task data leakage where a document seen during 
+    summarization training could boost topic classification on validation/test.
+    
+    Args:
+        primary_examples: Topic examples to filter
+        reference_examples: Summarization examples to check against
+        split: Name of split being processed (for logging)
+    
+    Returns:
+        Tuple of (filtered_examples, num_removed)
+    """
+    ref_fps = {_text_fingerprint(ex.source) for ex in reference_examples}
+    
+    filtered = []
+    removed = 0
+    for ex in primary_examples:
+        fp = _text_fingerprint(ex.text)
+        if fp in ref_fps:
+            removed += 1
+        else:
+            filtered.append(ex)
+    
+    if removed > 0:
+        print(f"  Dedup: removed {removed} overlapping examples from topic {split}")
+    
+    return filtered, removed

src/models/factory.py

@@ -548,13 +548,15 @@ def build_multitask_model(
         "summarization",
         LMHead(d_model=cfg.d_model, vocab_size=vocab_size, tie_embedding=decoder.embedding),
     )
-    # Emotion head with 2-layer MLP for better multi-label capacity (28 classes)
+    # Emotion head with attention pooling + 2-layer MLP for better multi-label capacity (28 classes)
+    # Attention pooling is superior to mean pooling for encoder-decoder models where
+    # hidden states are optimized for cross-attention rather than simple averaging.
     model.add_head(
         "emotion",
         ClassificationHead(
             d_model=cfg.d_model, 
             num_labels=num_emotions, 
-            pooler="mean", 
+            pooler="attention", 
             dropout=cfg.dropout,
             hidden_dim=cfg.d_model // 2,  # 384-dim hidden layer
         ),

src/models/heads.py

@@ -1,7 +1,7 @@
 """Prediction heads for Transformer models.
 
 This module provides task-specific output heads:
-- ClassificationHead: Sequence-level classification with pooling (mean/cls/max)
+- ClassificationHead: Sequence-level classification with pooling (mean/cls/max/attention)
 - TokenClassificationHead: Per-token classification (NER, POS tagging)
 - LMHead: Language modeling logits with optional weight tying
 - ProjectionHead: MLP for representation learning / contrastive tasks
@@ -14,6 +14,35 @@ from typing import Literal, Optional
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
+
+
+class AttentionPooling(nn.Module):
+    """Learned attention pooling over sequence positions.
+
+    Computes a weighted sum of hidden states using a learned query vector,
+    producing a single fixed-size representation. This is generally superior
+    to mean pooling for classification tasks on encoder-decoder models where
+    hidden states are optimized for cross-attention rather than pooling.
+    """
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.query = nn.Linear(d_model, 1, bias=False)
+
+    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        x: (batch, seq_len, d_model)
+        mask: (batch, seq_len) - True for valid tokens, False for padding
+        returns: (batch, d_model)
+        """
+        # Compute attention scores: (batch, seq_len, 1)
+        scores = self.query(x)
+        if mask is not None:
+            scores = scores.masked_fill(~mask.unsqueeze(-1), float("-inf"))
+        weights = F.softmax(scores, dim=1)  # (batch, seq_len, 1)
+        # Weighted sum: (batch, d_model)
+        return (weights * x).sum(dim=1)
 
 
 class ClassificationHead(nn.Module):
@@ -23,7 +52,7 @@ class ClassificationHead(nn.Module):
     Args:
         d_model: hidden size from encoder/decoder
         num_labels: number of output classes
-        pooler: one of 'mean', 'cls', 'max' - how to pool the sequence
+        pooler: one of 'mean', 'cls', 'max', 'attention' - how to pool the sequence
         dropout: dropout probability before final linear layer
         hidden_dim: optional intermediate dimension for 2-layer MLP (improves capacity)
     """
@@ -32,14 +61,17 @@ class ClassificationHead(nn.Module):
         self,
         d_model: int,
         num_labels: int,
-        pooler: Literal["mean", "cls", "max"] = "mean",
+        pooler: Literal["mean", "cls", "max", "attention"] = "mean",
         dropout: float = 0.1,
         hidden_dim: Optional[int] = None,
     ):
         super().__init__()
-        assert pooler in ("mean", "cls", "max"), "pooler must be 'mean'|'cls'|'max'"
+        assert pooler in ("mean", "cls", "max", "attention"), "pooler must be 'mean'|'cls'|'max'|'attention'"
         self.pooler = pooler
         self.dropout = nn.Dropout(dropout)
+
+        if pooler == "attention":
+            self.attn_pool = AttentionPooling(d_model)
         
         # Optional 2-layer MLP for more capacity (useful for multi-label)
         if hidden_dim is not None:
@@ -58,19 +90,14 @@ class ClassificationHead(nn.Module):
         mask: (batch, seq_len) - True for valid tokens, False for padding
         returns: (batch, num_labels)
         """
-        if self.pooler == "mean":
+        if self.pooler == "attention":
+            pooled = self.attn_pool(x, mask)
+        elif self.pooler == "mean":
             if mask is not None:
-                # mask is (B, S)
-                # x is (B, S, D)
-                # Expand mask to (B, S, 1)
                 mask_expanded = mask.unsqueeze(-1).float()
-                # Zero out padding
                 x = x * mask_expanded
-                # Sum over sequence
                 sum_embeddings = x.sum(dim=1)
-                # Count valid tokens
                 sum_mask = mask_expanded.sum(dim=1)
-                # Avoid division by zero
                 sum_mask = torch.clamp(sum_mask, min=1e-9)
                 pooled = sum_embeddings / sum_mask
             else:
@@ -79,7 +106,6 @@ class ClassificationHead(nn.Module):
             pooled = x[:, 0, :]
         else:  # max
             if mask is not None:
-                # Mask padding with -inf
                 mask_expanded = mask.unsqueeze(-1)
                 x = x.masked_fill(~mask_expanded, float("-inf"))
             pooled, _ = x.max(dim=1)

src/training/metrics.py

@@ -110,9 +110,9 @@ def calculate_bertscore(
     )
     
     return {
-        "precision": float(P.mean().item()),
-        "recall": float(R.mean().item()),
-        "f1": float(F1.mean().item()),
+        "precision": float(P.mean().item()),  # type: ignore[union-attr]
+        "recall": float(R.mean().item()),  # type: ignore[union-attr]
+        "f1": float(F1.mean().item()),  # type: ignore[union-attr]
     }
 
 
@@ -239,3 +239,213 @@ def get_confusion_matrix(
 ) -> np.ndarray:
     """Compute confusion matrix."""
     return cast(np.ndarray, confusion_matrix(targets, predictions, labels=labels))
+
+
+# --------------- Multi-label Emotion Metrics ---------------
+
+
+def multilabel_macro_f1(predictions: torch.Tensor, targets: torch.Tensor) -> float:
+    """Compute macro F1: average F1 per class (as in GoEmotions paper).
+    
+    This averages F1 across labels, giving equal weight to each emotion class 
+    regardless of prevalence. Directly comparable to GoEmotions baselines.
+    """
+    preds = predictions.float()
+    gold = targets.float()
+    
+    # Per-class TP, FP, FN
+    tp = (preds * gold).sum(dim=0)
+    fp = (preds * (1 - gold)).sum(dim=0)
+    fn = ((1 - preds) * gold).sum(dim=0)
+    
+    precision = tp / (tp + fp).clamp(min=1e-8)
+    recall = tp / (tp + fn).clamp(min=1e-8)
+    f1 = (2 * precision * recall) / (precision + recall).clamp(min=1e-8)
+    
+    # Zero out F1 for classes with no support in either predictions or targets
+    mask = (tp + fp + fn) > 0
+    if mask.sum() == 0:
+        return 0.0
+    return float(f1[mask].mean().item())
+
+
+def multilabel_micro_f1(predictions: torch.Tensor, targets: torch.Tensor) -> float:
+    """Compute micro F1: aggregate TP/FP/FN across all classes.
+    
+    This gives more weight to frequent classes. Useful when class distribution matters.
+    """
+    preds = predictions.float()
+    gold = targets.float()
+    
+    tp = (preds * gold).sum()
+    fp = (preds * (1 - gold)).sum()
+    fn = ((1 - preds) * gold).sum()
+    
+    precision = tp / (tp + fp).clamp(min=1e-8)
+    recall = tp / (tp + fn).clamp(min=1e-8)
+    f1 = (2 * precision * recall) / (precision + recall).clamp(min=1e-8)
+    return float(f1.item())
+
+
+def multilabel_per_class_metrics(
+    predictions: torch.Tensor,
+    targets: torch.Tensor,
+    class_names: Sequence[str] | None = None,
+) -> Dict[str, Dict[str, float]]:
+    """Compute per-class precision, recall, F1 for multi-label classification.
+    
+    Returns a dict mapping class name/index to its metrics.
+    """
+    preds = predictions.float()
+    gold = targets.float()
+    num_classes = preds.shape[1]
+    
+    tp = (preds * gold).sum(dim=0)
+    fp = (preds * (1 - gold)).sum(dim=0)
+    fn = ((1 - preds) * gold).sum(dim=0)
+    
+    report: Dict[str, Dict[str, float]] = {}
+    for i in range(num_classes):
+        name = class_names[i] if class_names else str(i)
+        p = (tp[i] / (tp[i] + fp[i]).clamp(min=1e-8)).item()
+        r = (tp[i] / (tp[i] + fn[i]).clamp(min=1e-8)).item()
+        f = (2 * p * r) / max(p + r, 1e-8)
+        report[name] = {
+            "precision": p,
+            "recall": r,
+            "f1": f,
+            "support": int(gold[:, i].sum().item()),
+        }
+    return report
+
+
+def tune_per_class_thresholds(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    thresholds: Sequence[float] | None = None,
+) -> tuple[List[float], float]:
+    """Tune per-class thresholds on validation set to maximize macro F1.
+    
+    For each class, tries multiple thresholds and selects the one that 
+    maximizes that class's F1 score. This is standard practice for multi-label 
+    classification (used in the original GoEmotions paper).
+    
+    Args:
+        logits: Raw model logits (batch, num_classes)
+        targets: Binary target labels (batch, num_classes)
+        thresholds: Candidate thresholds to try (default: 0.1 to 0.9 by 0.05)
+    
+    Returns:
+        Tuple of (best_thresholds_per_class, resulting_macro_f1)
+    """
+    if thresholds is None:
+        thresholds = [round(t, 2) for t in np.arange(0.1, 0.9, 0.05).tolist()]
+    
+    probs = torch.sigmoid(logits)
+    num_classes = probs.shape[1]
+    gold = targets.float()
+    
+    best_thresholds: List[float] = []
+    for c in range(num_classes):
+        best_f1 = -1.0
+        best_t = 0.5
+        for t in thresholds:
+            preds = (probs[:, c] >= t).float()
+            tp = (preds * gold[:, c]).sum()
+            fp = (preds * (1 - gold[:, c])).sum()
+            fn = ((1 - preds) * gold[:, c]).sum()
+            if tp + fp > 0 and tp + fn > 0:
+                p = tp / (tp + fp)
+                r = tp / (tp + fn)
+                f1 = (2 * p * r / (p + r)).item()
+            else:
+                f1 = 0.0
+            if f1 > best_f1:
+                best_f1 = f1
+                best_t = t
+        best_thresholds.append(best_t)
+    
+    # Compute resulting macro F1 with tuned thresholds
+    tuned_preds = torch.zeros_like(probs)
+    for c in range(num_classes):
+        tuned_preds[:, c] = (probs[:, c] >= best_thresholds[c]).float()
+    macro_f1 = multilabel_macro_f1(tuned_preds, targets)
+    
+    return best_thresholds, macro_f1
+
+
+# --------------- Statistical Tests ---------------
+
+
+def bootstrap_confidence_interval(
+    scores: Sequence[float],
+    n_bootstrap: int = 1000,
+    confidence: float = 0.95,
+    seed: int = 42,
+) -> tuple[float, float, float]:
+    """Compute bootstrap confidence interval for a metric.
+    
+    Args:
+        scores: Per-sample metric values
+        n_bootstrap: Number of bootstrap resamples
+        confidence: Confidence level (default 95%)
+        seed: Random seed for reproducibility
+    
+    Returns:
+        Tuple of (mean, lower_bound, upper_bound)
+    """
+    rng = np.random.default_rng(seed)
+    scores_arr = np.array(scores)
+    n = len(scores_arr)
+    
+    bootstrap_means = []
+    for _ in range(n_bootstrap):
+        sample = rng.choice(scores_arr, size=n, replace=True)
+        bootstrap_means.append(float(np.mean(sample)))
+    
+    bootstrap_means.sort()
+    alpha = 1 - confidence
+    lower_idx = int(alpha / 2 * n_bootstrap)
+    upper_idx = int((1 - alpha / 2) * n_bootstrap)
+    
+    return (
+        float(np.mean(scores_arr)),
+        bootstrap_means[lower_idx],
+        bootstrap_means[min(upper_idx, n_bootstrap - 1)],
+    )
+
+
+def paired_bootstrap_test(
+    scores_a: Sequence[float],
+    scores_b: Sequence[float],
+    n_bootstrap: int = 10000,
+    seed: int = 42,
+) -> float:
+    """Paired bootstrap significance test between two systems.
+    
+    Tests if system B is significantly better than system A.
+    
+    Args:
+        scores_a: Per-sample scores from system A
+        scores_b: Per-sample scores from system B
+        n_bootstrap: Number of bootstrap iterations
+        seed: Random seed
+    
+    Returns:
+        p-value (probability that B is not better than A)
+    """
+    rng = np.random.default_rng(seed)
+    a = np.array(scores_a)
+    b = np.array(scores_b)
+    assert len(a) == len(b), "Both score lists must have the same length"
+    
+    n = len(a)
+    
+    count = 0
+    for _ in range(n_bootstrap):
+        idx = rng.choice(n, size=n, replace=True)
+        diff = float(np.mean(b[idx]) - np.mean(a[idx]))
+        if diff <= 0:
+            count += 1
+    
+    return count / n_bootstrap

src/training/trainer.py

@@ -7,6 +7,8 @@ Handles training across summarization, emotion, and topic heads with:
 - Cosine LR schedule with warmup
 - Early stopping
 - MLflow logging
+- Temperature-based task sampling (configurable alpha)
+- Gradient conflict diagnostics
 
 Author: Oliver Perrin
 Date: December 2025
@@ -22,6 +24,7 @@ from dataclasses import dataclass
 from typing import Any, Callable, Dict, List
 
 import mlflow
+import numpy as np
 import torch
 import torch.nn.functional as F
 from torch.optim.lr_scheduler import LambdaLR
@@ -53,6 +56,16 @@ class TrainerConfig:
     # Early stopping
     early_stopping_patience: int | None = 5
     
+    # Task sampling strategy: "round_robin" or "temperature"
+    # Temperature sampling: p_i ∝ n_i^alpha where n_i = dataset size
+    # alpha < 1 reduces dominance of large tasks (recommended: 0.5-0.7)
+    task_sampling: str = "round_robin"
+    task_sampling_alpha: float = 0.5
+    
+    # Gradient conflict diagnostics
+    # Compute inter-task gradient cosine similarity every N steps (0 = disabled)
+    gradient_conflict_frequency: int = 0
+    
     # MLflow
     experiment_name: str = "LexiMind"
     run_name: str | None = None
@@ -210,7 +223,7 @@ class Trainer:
         train: bool,
         epoch: int,
     ) -> Dict[str, float]:
-        """Run one epoch."""
+        """Run one epoch with configurable task sampling strategy."""
         self.model.train(train)
         metrics: Dict[str, List[float]] = defaultdict(list)
         iterators = {task: iter(loader) for task, loader in loaders.items()}
@@ -220,12 +233,33 @@ class Trainer:
         phase = "Train" if train else "Val"
         pbar = tqdm(range(max_batches), desc=f"  {phase}", leave=False, file=sys.stderr)
 
+        # Temperature-based task sampling: p_i ∝ n_i^alpha
+        task_names = list(loaders.keys())
+        if self.config.task_sampling == "temperature" and len(task_names) > 1:
+            sizes = np.array([len(loaders[t].dataset) for t in task_names], dtype=np.float64)  # type: ignore[arg-type]
+            alpha = self.config.task_sampling_alpha
+            probs = sizes ** alpha
+            probs = probs / probs.sum()
+            tqdm.write(f"  Temperature sampling (α={alpha}): " +
+                       ", ".join(f"{t}={p:.2%}" for t, p in zip(task_names, probs, strict=True)))
+        else:
+            probs = None
+
         ctx = torch.enable_grad() if train else torch.no_grad()
         with ctx:
             for step in pbar:
                 step_loss = 0.0
 
-                for task, loader in loaders.items():
+                # Select tasks for this step
+                if probs is not None and train:
+                    # Temperature sampling: sample tasks based on dataset size
+                    selected_tasks = list(np.random.choice(task_names, size=len(task_names), replace=True, p=probs))
+                else:
+                    # Round-robin: all tasks every step
+                    selected_tasks = task_names
+
+                for task in selected_tasks:
+                    loader = loaders[task]
                     batch = self._get_batch(iterators, loader, task)
                     if batch is None:
                         continue
@@ -253,6 +287,14 @@ class Trainer:
                         scaled = (loss * weight) / accum
                         scaled.backward()
 
+                # Gradient conflict diagnostics
+                if (train and self.config.gradient_conflict_frequency > 0
+                        and (step + 1) % self.config.gradient_conflict_frequency == 0):
+                    conflict_stats = self._compute_gradient_conflicts(loaders, iterators)
+                    for k, v in conflict_stats.items():
+                        metrics[f"grad_{k}"].append(v)
+                        mlflow.log_metric(f"grad_{k}", v, step=self.global_step)
+
                 # Optimizer step
                 if train and (step + 1) % accum == 0:
                     torch.nn.utils.clip_grad_norm_(
@@ -415,6 +457,56 @@ class Trainer:
         tqdm.write(f"{'=' * 50}\n")
         self.model.train()
 
+    def _compute_gradient_conflicts(
+        self,
+        loaders: Dict[str, DataLoader],
+        iterators: Dict,
+    ) -> Dict[str, float]:
+        """Compute inter-task gradient cosine similarity to diagnose conflicts.
+        
+        Returns cosine similarity between gradient vectors for each task pair.
+        Negative values indicate conflicting gradients (negative transfer risk).
+        """
+        task_grads: Dict[str, torch.Tensor] = {}
+        
+        for task, loader in loaders.items():
+            self.optimizer.zero_grad()
+            batch = self._get_batch(iterators, loader, task)
+            if batch is None:
+                continue
+            
+            dtype = torch.bfloat16 if self.use_bfloat16 else torch.float16
+            with torch.autocast("cuda", dtype=dtype, enabled=self.use_amp):
+                loss, _ = self._forward_task(task, batch)
+            
+            if torch.isnan(loss):
+                continue
+            
+            loss.backward()
+            
+            # Flatten all gradients into a single vector
+            grad_vec = []
+            for p in self.model.parameters():
+                if p.grad is not None:
+                    grad_vec.append(p.grad.detach().clone().flatten())
+            if grad_vec:
+                task_grads[task] = torch.cat(grad_vec)
+        
+        self.optimizer.zero_grad()
+        
+        # Compute pairwise cosine similarity
+        stats: Dict[str, float] = {}
+        tasks = list(task_grads.keys())
+        for i in range(len(tasks)):
+            for j in range(i + 1, len(tasks)):
+                t1, t2 = tasks[i], tasks[j]
+                g1, g2 = task_grads[t1], task_grads[t2]
+                cos_sim = F.cosine_similarity(g1.unsqueeze(0), g2.unsqueeze(0)).item()
+                stats[f"cos_sim_{t1}_{t2}"] = cos_sim
+                stats[f"conflict_{t1}_{t2}"] = 1.0 if cos_sim < 0 else 0.0
+        
+        return stats
+
     def _log_config(self) -> None:
         """Log config to MLflow."""
         mlflow.log_params({