Multi-Task Representation Learning for Conservative Linear Bandits

Abstract: This paper presents the Constrained Multi-Task Representation Learning (CMTRL) framework for linear bandits. We consider T linear bandit tasks in a d dimensional space, which share a common low-dimensional representation of dimension r, where r is much smaller than the minimum of d and T. Furthermore, tasks are constrained so that only actions meeting specific safety or performance requirements are allowed, referred to as conservative (safe) bandits. We introduce a novel algorithm, Safe-Alternating projected Gradient Descent and minimization (Safe-AltGDmin), to recover a low-rank feature matrix while satisfying the given constraints. Building on this algorithm, we propose a multi-task representation learning framework for conservative linear bandits and establish theoretical guarantees for its regret and sample complexity bounds. We presented experiments and compared the performance of our algorithm with benchmark algorithms.

• The paper introduces Safe-AltGDmin, which combines conservative exploration with alternating gradient descent to enforce strict safety constraints in multi-task linear bandits.
• It leverages a shared low-rank representation to achieve exponential error decay and significant sample efficiency compared to independent task estimation.
• Empirical results validate its lower cumulative regret and robust performance on synthetic and real datasets, demonstrating scalability in high-dimensional settings.

Multi-Task Representation Learning under Safety Constraints in Linear Bandits

Context and Motivation

The paper "Multi-Task Representation Learning for Conservative Linear Bandits" (2605.12176) addresses a significant gap in sequential decision-making research: multi-task learning with explicit safety or performance constraints in linear bandit settings. Classic multi-task representation learning (MTRL) leverages shared structure among related tasks—a common low-dimensional subspace—to improve data efficiency and reduce regret in bandits. However, the intersection of MTRL and conservative (safe) bandits remains largely unexplored, despite safety constraints being central in applications such as healthcare policy design, financial portfolio management, and autonomous systems, where regulatory, ethical, or operational boundaries are present.

Problem Formulation

The authors formalize the Constrained Multi-Task Representation Learning (CMTRL) problem as follows:

• $T$ linear bandit tasks, each with unknown parameter $\theta_t \in \mathbb{R}^d$, are assumed to share a low-rank structure: the $$\Theta^* = [\theta_1^*, ..., \theta_T^*] \in \mathbb{R}^{d \times T}$$ matrix is rank-$r$ ($r \ll \min\{d, T\}$), factorized as $\Theta^* = B^*W^*$.
• At each round $n \in [N]$, each task independently chooses an action $x_{n,t}$, observes reward $y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t}$, and must ensure $x_{n,t}^\top \theta_t^* \geq (1-\alpha) r_{b,n,t}$ where $\theta_t \in \mathbb{R}^d$0 is a baseline reward and $\theta_t \in \mathbb{R}^d$1 determines conservatism.
• All actions throughout the learning horizon must satisfy the above stage-wise constraint.

Unlike previous works on unconstrained MTRL in bandits or conservative single-task settings, this work analyzes both representation sharing and global, per-epoch, or per-round safety constraints.

Algorithmic Contribution: Safe-AltGDmin

The authors propose Safe-AltGDmin, a two-phase, epoch-based algorithm comprising:

1. Safe (Conservative) Exploration: In epoch-1, a parameter $\theta_t \in \mathbb{R}^d$2 interpolates between the strictly safe baseline action and independent random exploration, controlling trade-off between safety and exploration. Parameter $\theta_t \in \mathbb{R}^d$3 is upper bounded analytically to ensure safety with high probability via Gaussian concentration bounds.
2. Alternating Gradient Descent and Minimization (AltGDmin): At the end of each epoch, the method alternates between gradient steps (for the common feature extractor $\theta_t \in \mathbb{R}^d$4) and closed-form minimization (for the task coefficients $\theta_t \in \mathbb{R}^d$5) to minimize the empirical squared error over all collected samples. A robust spectral initialization (with statistical truncation to avoid outliers) is used to avoid poor local minima due to non-convexity.
3. Epoch-wise Greedy Exploration: In subsequent epochs, actions are chosen greedily with respect to the latest parameter estimates subject to the safety constraint, and AltGDmin steps update the shared representation and task parameters.

These phases are recursively interleaved, and parameter settings for stepsizes, number of gradient steps, and conservative interpolation are derived to guarantee safety and convergence.

Theoretical Guarantees

The authors' analysis is exacting:

• Safety Guarantee: For all rounds and epochs, with high probability, every selected action satisfies the stage-wise conservative constraint. This guarantee is nontrivial under simultaneous estimation and safety, especially in the multi-task, non-convex, representation learning setting.
• Subspace Recovery: Under modest incoherence and i.i.d. Gaussian design/noise assumptions, the subspace distance between the estimated and true feature extractors decreases exponentially fast with the number of AltGDmin iterations. The contraction factor is explicit.
• Sample Complexity: To achieve target estimation error $\theta_t \in \mathbb{R}^d$6 and robust safety, the total sample complexity across $\theta_t \in \mathbb{R}^d$7 tasks per epoch scales as $\theta_t \in \mathbb{R}^d$8, where NSR is the noise-to-signal ratio. This is dramatically lower than $\theta_t \in \mathbb{R}^d$9 samples required for independent per-task estimation.
• Regret Bound: The cumulative regret $$\Theta^* = [\theta_1^*, ..., \theta_T^*] \in \mathbb{R}^{d \times T}$$0 is shown to obey $$\Theta^* = [\theta_1^*, ..., \theta_T^*] \in \mathbb{R}^{d \times T}$$1, demonstrating sublinear scaling in both number of tasks $$\Theta^* = [\theta_1^*, ..., \theta_T^*] \in \mathbb{R}^{d \times T}$$2 and horizon $$\Theta^* = [\theta_1^*, ..., \theta_T^*] \in \mathbb{R}^{d \times T}$$3 (up to log factors and problem constants). This is a sharp improvement over $$\Theta^* = [\theta_1^*, ..., \theta_T^*] \in \mathbb{R}^{d \times T}$$4 for single-task, constraint-unaware methods.
• Communication/Time Complexity: The authors derive tight bounds reflecting AltGDmin's scalability.

Notably, the regret guarantee does not assume exact knowledge of the underlying representation, and the algorithm tolerates both statistical and optimization error, even under the imposed constraints.

Empirical Results

The proposed algorithm is comprehensively evaluated on:

• Synthetic data with controllable $$\Theta^* = [\theta_1^*, ..., \theta_T^*] \in \mathbb{R}^{d \times T}$$5, $$\Theta^* = [\theta_1^*, ..., \theta_T^*] \in \mathbb{R}^{d \times T}$$6, $$\Theta^* = [\theta_1^*, ..., \theta_T^*] \in \mathbb{R}^{d \times T}$$7 parameters.
• The Movielens-100K recommendation benchmark (processed for the linear contextual bandit model).

Key empirical results:

• Constraint Satisfaction: Safe-AltGDmin strictly enforces the per-round constraint (i.e., zero empirical violation) across all datasets and regime choices. Competitors which ignore constraints incur frequent violations.
• Regret: Safe-AltGDmin achieves lower cumulative regret than the following baselines: trace norm convex relaxation MTRL, Thompson sampling for independent safe single-tasks, and the method-of-moments (MoM) MTRL approach. The performance gap increases with dimensionality and number of tasks.
• Estimation Error: The algorithm maintains a lower parameter estimation error—especially as the number of tasks increases—validating the theoretical scaling benefits of leveraging shared low-rank representations.
• Robustness to Data Scarcity: The method capitalizes on information sharing across tasks, especially in data-scarce regimes, highlighting a practical advantage where safety exploration significantly limits feasible trajectories.

Implications and Future Directions

From a theoretical perspective, this work extends the frontiers of representation-based learning in bandits by ensuring sustained constraint satisfaction across all exploration-exploitation phases, with formal convergence and regret guarantees. It handles the intractability of non-convex optimization via provable, robust initialization and sample-efficient AltGDmin updates.

Practically, Safe-AltGDmin is directly relevant for applications demanding reliable sequential decision making with hard constraints—e.g., clinical trials, automated trading, and deployment of recommender systems with fairness or fiduciary restrictions.

Future research could address several open questions:

• Relaxation of Gaussian design and noise assumptions prevalent in current analysis.
• Extending theoretical guarantees to adaptive, possibly adversarial, safety constraints or reward drifts.
• Merging nonlinearity (e.g., kernelized or neural MTRL) with the constraint framework to capture broader real-world phenomena.
• Federated or distributed implementations where safety/communication trade-offs may be even more critical.

Conclusion

This paper establishes a rigorous foundation for multi-task, constraint-aware representation learning in linear bandits. The Safe-AltGDmin algorithm achieves both provably safe exploration and efficient feature sharing, yielding exponential estimation error decay and sublinear cumulative regret with respect to the number of tasks and horizon. Empirical results confirm performance benefits over prior benchmarks. The work opens several directions at the intersection of online learning, low-rank modeling, and safety-critical sequential decision-making.