Safe-AltGDmin: Constrained Multi-Task Bandits

• The paper introduces Safe-AltGDmin, a safe alternating projected gradient descent minimization procedure for constrained multi-task representation learning in linear bandit settings.
• It employs spectral initialization, alternating least-squares updates, and sample splitting to achieve exponential subspace contraction and sublinear regret.
• The method integrates conservative safe exploration with rigorous theoretical guarantees and empirical validations on synthetic data and Movielens-100K.

Searching arXiv for the cited papers and related work to ground the article. Safe-AltGDmin is the name given to the “Safe-Alternating projected Gradient Descent and minimization” procedure introduced for Constrained Multi-Task Representation Learning (CMTRL) in conservative linear bandits. It addresses a multi-task setting in which task parameters share a common low-dimensional representation, while every action must satisfy a stage-wise performance constraint relative to a known baseline. The method combines conservative exploration in the first epoch, truncated spectral initialization, alternating least-squares updates for task-specific coefficients, projected gradient descent on a shared feature extractor, and sample splitting to obtain high-probability safety, exponential subspace contraction, and sublinear regret (Lin et al., 12 May 2026).

1. Problem setting and formal objective

Safe-AltGDmin is formulated for $T$ linear bandit tasks in a $d$ dimensional space. At round $n \in \{1,\dots,N\}$, task $t$ selects an action $x_{n,t} \in \mathbb{R}^d$ and observes

$y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t},$

where $\theta_t^* \in \mathbb{R}^d$ is unknown and the noise variables are i.i.d. zero-mean sub-Gaussian, with the analysis using Gaussian noise of variance $\sigma^2$. The tasks share an $r$-dimensional representation, with $r \ll \min\{d,T\}$, so that there exist $d$0 with orthonormal columns and $d$1 such that

$d$2

The singular spectrum of $d$3 determines $d$4, $d$5, the condition number $d$6, and the noise-to-signal ratio $d$7 (Lin et al., 12 May 2026).

The defining safety requirement is stage-wise and baseline-relative. For each task $d$8 and round $d$9, a baseline action $n \in \{1,\dots,N\}$0 and its known baseline reward $n \in \{1,\dots,N\}$1 are given, and the learner must satisfy

$n \in \{1,\dots,N\}$2

with known $n \in \{1,\dots,N\}$3. The safe set for task $n \in \{1,\dots,N\}$4 is therefore

$n \in \{1,\dots,N\}$5

Because $n \in \{1,\dots,N\}$6 is unknown, safety is enforced indirectly: epoch $n \in \{1,\dots,N\}$7 uses conservative randomized mixtures of the baseline, and later epochs use greedy actions with respect to the current estimate, together with a high-probability safety proof (Lin et al., 12 May 2026).

The estimation problem is organized by epochs. If the horizon is partitioned as $n \in \{1,\dots,N\}$8, then in epoch $n \in \{1,\dots,N\}$9 the factorization $t$0 is fitted by minimizing the nonconvex squared loss

$t$1

optionally with the Stiefel constraint $t$2, which the algorithm enforces through a projection or QR step (Lin et al., 12 May 2026).

2. Alternating projected GD and minimization

The algorithm alternates between minimization in the task-specific block $t$3 and projected gradient descent in the shared block $t$4. This is the same structural pattern as the broader AltGDmin framework for partly-decoupled optimization, in which the loss is differentiable in one block and the other block is fast to minimize, often because it decouples across clients or samples (Vaswani, 20 Apr 2025).

In Safe-AltGDmin, epoch $t$5 is a safe exploration phase. For each task $t$6 and each round $t$7, the action is

$t$8

where the user parameter $t$9 trades exploration and safety. The paper also requires $x_{n,t} \in \mathbb{R}^d$0 for contraction in epoch $x_{n,t} \in \mathbb{R}^d$1, so that the Gram matrix remains well-conditioned (Lin et al., 12 May 2026).

Initialization is spectral and truncation-based. Let $x_{n,t} \in \mathbb{R}^d$2 be the $x_{n,t} \in \mathbb{R}^d$3 design matrix of epoch-$x_{n,t} \in \mathbb{R}^d$4 contexts for task $x_{n,t} \in \mathbb{R}^d$5, and let $x_{n,t} \in \mathbb{R}^d$6 collect the corresponding rewards. The unnormalized method-of-moments estimator is

$x_{n,t} \in \mathbb{R}^d$7

To mitigate heavy tails, the labels are truncated entrywise using a global threshold $x_{n,t} \in \mathbb{R}^d$8, giving $x_{n,t} \in \mathbb{R}^d$9, and then

$y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t},$0

The initial factor $y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t},$1 is set to the top-$y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t},$2 left singular vectors of $y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t},$3 (Lin et al., 12 May 2026).

Within each epoch, the data are partitioned into $y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t},$4 disjoint blocks per task. This sample-splitting design removes statistical dependencies between the $y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t},$5-update and the $y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t},$6-gradient computation. For $y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t},$7, the algorithm executes two steps. First, for each task $y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t},$8, it solves ordinary least squares on block $y_{n,t} = x_{n,t}^\top \theta_t^* + \eta_{n,t},$9:

$\theta_t^* \in \mathbb{R}^d$0

This gives the closed-form update

$\theta_t^* \in \mathbb{R}^d$1

Second, using the independent block $\theta_t^* \in \mathbb{R}^d$2, it computes an empirical gradient of $\theta_t^* \in \mathbb{R}^d$3 at $\theta_t^* \in \mathbb{R}^d$4,

$\theta_t^* \in \mathbb{R}^d$5

takes the step

$\theta_t^* \in \mathbb{R}^d$6

and projects onto the Stiefel set

$\theta_t^* \in \mathbb{R}^d$7

by QR or polar decomposition:

$\theta_t^* \in \mathbb{R}^d$8

After $\theta_t^* \in \mathbb{R}^d$9 such iterations, the epoch output is $\sigma^2$0, $\sigma^2$1, and $\sigma^2$2 (Lin et al., 12 May 2026).

3. Safety mechanisms and the meaning of “safe”

The “safe” component of Safe-AltGDmin is algorithmic and statistical rather than a separate projection onto an explicit safe action set. In epoch $\sigma^2$3, safety is guaranteed by choosing $\sigma^2$4 small enough. A sufficient condition states that, for confidence $\sigma^2$5,

$\sigma^2$6

which yields, with probability at least $\sigma^2$7,

$\sigma^2$8

for all $\sigma^2$9 and all tasks $r$0. Since $r$1 is unknown, the analysis replaces it by the computable upper bound $r$2 derived from the incoherence parameter $r$3 (Lin et al., 12 May 2026).

In later epochs, the algorithm uses the greedy action

$r$4

without an explicit safe-set restriction. Safety is instead proved by combining a uniform Gaussian concentration bound for

$r$5

with exponential contraction of the estimation error. Under the sample-size and iteration conditions of the main theorem, every greedy action in epochs $r$6 satisfies

$r$7

with high probability (Lin et al., 12 May 2026).

Several additional devices stabilize the optimization itself. The step size is chosen small enough that $r$8 is positive semidefinite; a practical form consistent with the proofs is

$r$9

The only explicit constraint set projected onto in the CMTRL algorithm is the Stiefel set $r \ll \min\{d,T\}$0, while $r \ll \min\{d,T\}$1 remains unconstrained. This aligns Safe-AltGDmin with the broader AltGDmin literature, where projections, safe constant step sizes, robust spectral initializers, and sample splitting are the primary stability mechanisms rather than line search or trust-region procedures (Lin et al., 12 May 2026, Vaswani, 20 Apr 2025).

4. Assumptions and theoretical guarantees

The analysis assumes a shared representation $r \ll \min\{d,T\}$2 with orthonormal $r \ll \min\{d,T\}$3, bounded baseline gaps $r \ll \min\{d,T\}$4, Gaussian design and noise in epoch $r \ll \min\{d,T\}$5, and column-wise incoherence of $r \ll \min\{d,T\}$6:

$r \ll \min\{d,T\}$7

These assumptions ensure that no single task dominates and that the effective local quadratic model has the strong convexity and smoothness needed for contraction (Lin et al., 12 May 2026).

The initialization guarantee states that if

$r \ll \min\{d,T\}$8

then, with probability at least $r \ll \min\{d,T\}$9,

$d$00

where

$d$01

In epoch $d$02, if $d$03, if $d$04, and if the prescribed sample-size conditions hold, then with probability at least $d$05,

$d$06

For later epochs, if $d$07 is sufficiently small and the per-epoch sample sizes are large enough at target accuracy $d$08, then

$d$09

with probability at least $d$10, yielding exponential decay down to $d$11 (Lin et al., 12 May 2026).

The task-parameter errors inherit the same contraction:

$d$12

The safety theorem combines these estimation bounds with lower bounds on the per-epoch sample sizes and the requirement

$d$13

or the explicit logarithmic form appearing in the proof, to show that every greedy action in epochs $d$14 is stage-wise safe with high probability (Lin et al., 12 May 2026).

The regret guarantee is

$d$15

with probability at least $d$16. Ignoring constants and logarithms, this gives

$d$17

which depends on the intrinsic rank $d$18 rather than the ambient dimension $d$19 (Lin et al., 12 May 2026).

5. Relation to AltGDmin and federated low-rank recovery

Safe-AltGDmin is a specialized member of the AltGDmin family. In the general partly-decoupled template, one solves

$d$20

where $d$21 is differentiable in $d$22 and the minimization over $d$23 is closed-form or reliably solved. When

$d$24

the $d$25-update decomposes into $d$26 local problems, and clients transmit only partial gradients in the shared block $d$27, making AltGDmin communication-efficient in vertically federated settings (Vaswani, 20 Apr 2025).

The low-rank recovery literature that preceded Safe-AltGDmin already contained most of the optimization safeguards later emphasized in the bandit setting. For low-rank column-wise sensing, AltGDmin alternates closed-form updates

$d$28

with a projected gradient step

$d$29

and uses truncated spectral initialization, sample splitting, and constant step sizes such as $d$30 to obtain geometric convergence in subspace distance (Vaswani, 2023). For federated low-rank matrix completion, the algorithm similarly alternates decoupled least squares over columns with a masked gradient step and QR on $d$31, while row-incoherence projection at initialization preserves the assumptions needed for concentration and contraction (Abbasi et al., 2024).

A recurrent misconception is that “Safe-AltGDmin” names a universally standardized variant across the entire AltGDmin literature. The record is more specific. The CMTRL bandit paper explicitly introduces “Safe-AltGDmin” as a named algorithm (Lin et al., 12 May 2026). By contrast, the general AltGDmin framework paper states that no named “Safe-AltGDmin” variant is introduced there, even though its core method already incorporates safety-critical mechanisms such as QR projection, explicit safe step-size choices, robust spectral initializers based on truncation, incoherence projection for LRMC initialization, and sample splitting (Vaswani, 20 Apr 2025). The LRMC paper likewise presents “Toward a ‘Safe-AltGDmin’” as a safety-enhanced variant grounded in the analysis rather than the paper’s primary named algorithm (Abbasi et al., 2024).

6. Empirical behavior, implementation choices, and limitations

The empirical study for Safe-AltGDmin uses both synthetic data and Movielens-100K. In the synthetic setup, the default parameters are $d$32, $d$33, ten actions per round per task, Gaussian noise variance $d$34, and a baseline equal to the 5th-best action; $d$35 and $d$36 are drawn from Gaussian distributions and orthonormalized. In the Movielens-100K setup, ratings are normalized to $d$37, matrix factorization is used, columns are clustered into $d$38 groups to define task-specific contexts, and the rank is set to $d$39. Baselines include a trace-norm convex relaxation, Thompson sampling with safe set estimation per task, and a method-of-moments estimator followed by least squares and greedy action selection (Lin et al., 12 May 2026).

The reported findings are threefold. First, the estimation error $d$40 is consistently lower than for the baselines and improves as $d$41 increases, reflecting the benefit of shared representation learning. Second, Safe-AltGDmin and Thompson sampling with safe set estimation incur near-zero safety violations, whereas trace-norm and method-of-moments baselines show many violations because they ignore constraints. Third, the regret displays sublinear growth in $d$42 and $d$43, consistent with the $d$44 theory; trace-norm and method-of-moments procedures can have lower regret only by violating safety (Lin et al., 12 May 2026).

The implementation guidance given for the algorithm is tightly coupled to the theory. The paper recommends choosing

$d$45

with $d$46, setting

$d$47

and taking

$d$48

or the explicit form from the safety theorem. The per-GD-iteration complexity is

$d$49

projection costs $d$50, total per-epoch complexity is

$d$51

and communication is $d$52 if the method is distributed (Lin et al., 12 May 2026).

The limitations stated for the method are equally specific. If tasks do not share a low-dimensional representation, so that the effective rank is large, the advantage diminishes because regret and sample complexity scale with $d$53. Tight safety parameters, meaning small $d$54, and weak baselines, meaning small $d$55, restrict epoch-$d$56 exploration because $d$57 must be very small. Context distributions that are adversarial or heavy-tailed beyond the truncation robustness require more advanced robust estimators. The paper identifies nonlinear shared representations, contextual safety constraints, adaptive rank selection, and distributed or federated implementations leveraging the same AltGDmin core as natural extensions (Lin et al., 12 May 2026).