BanditDB Docs — NeuralLinUCB

NeuralLinUCB

An MLP preprocessing stage that gives LinUCB a learnable, non-linear view of the context.

Standard LinUCB assumes the expected reward is a linear function of the raw context — which breaks down when features interact non-linearly or when the context vector carries redundant dimensions. NeuralLinUCB adds a small MLP that projects the caller's context into a compact fixed-size embedding; LinUCB then operates entirely in that embedding space. The network is retrained in batches on a background schedule — keeping the hot predict/reward path purely online with no ML overhead per request.

context x (context_dim floats) │ ┌────┴──────────────────────────────┐ │ hidden_layers × hidden_dim × ReLU│ MLP h(x ; W) └────┬──────────────────────────────┘ retrained in background, off the request path │ L2-normalise ▼ embedding h (embed_dim floats) │ θᵀh + α · √(hᵀ A⁻¹ h) ← LinUCB (Sherman-Morrison online update)

Two algorithms run in tandem:

Algorithm 1 (online, every request): Runs the same Sherman-Morrison LinUCB update used by "linucb", but on the MLP embedding rather than raw context. No added latency — the MLP forward pass is a single matrix multiply chain.
Algorithm 2 (batch, background): Once retrain_every rewards have accumulated, AdamW minimises ½·MSE(θᵀh(x;W), r) + (m·λ/2)·‖W − W₀‖² for retrain_steps steps. The L2 penalty toward the initialisation weights W₀ prevents catastrophic forgetting across retrain cycles. Arm matrices are then re-accumulated — replayed through the new embedding via Sherman-Morrison.

Thompson Sampling variant: "neural_thompson_sampling" uses the identical MLP embedding and Algorithm 2 batch retrain described above — the only difference is Algorithm 1. Instead of scoring by the UCB bound θᵀh + α·√(hᵀA⁻¹h), it draws w ~ N(θ, α²·A⁻¹) and scores by wᵀh. Same config schema, same table below. Tends toward better long-run convergence with higher early regret than NeuralLinUCB.

Configuration

Per-campaign, set at creation:

Parameter	Default	Meaning
`context_dim`	required	Length of the context vector the caller sends at predict time. Fixed at campaign creation.
`embed_dim`	32	Embedding (arm matrix) dimension. Arm matrices are `embed_dim × embed_dim`. Smaller = faster online updates; larger = more expressive. 16–64 covers most use cases.
`hidden_dim`	128	Width of each hidden layer.
`hidden_layers`	2	Number of hidden layers before the final embedding projection.
`retrain_every`	200	Rewards to accumulate before triggering Algorithm 2. Larger = fewer retrains, more stable arm matrices in between.
`retrain_steps`	100	AdamW gradient-descent steps per retrain cycle.
`learning_rate`	0.001	AdamW learning rate.
`lambda`	1.0	L2 regularization coefficient toward initial weights `W₀`. Higher = slower adaptation, more stable. Lower = allows faster weight drift between cycles.

Process-wide, set by environment variable:

Variable	Default	Meaning
`BANDITDB_RETRAIN_POLL_SECS`	`2`	How often the background worker checks whether any campaign is due for a retrain. `0` restores the pre-2.0.0 behaviour of retraining only at checkpoint.
`BANDITDB_NEURAL_BUFFER_CAP`	`50000`	Replay-buffer entries held per campaign. This bounds memory: entries × context_dim × 8 bytes, per campaign.
`BANDITDB_NEURAL_BATCH_SIZE`	`4000`	Minibatch sampled per gradient step. Decouples the cost of a step from buffer size, so a large buffer no longer makes every retrain proportionally slower.

db.create_campaign(
    "content",
    arms=["article_a", "article_b", "article_c"],
    algorithm={
        "neural_lin_ucb": {
            "context_dim": 32,    # must match the context vector you send at predict time
            "embed_dim": 16,      # arm matrices are 16 × 16
            "hidden_dim": 128,
            "hidden_layers": 2,
            "retrain_every": 500, # first retrain after 500 rewards
            "retrain_steps": 100,
            "learning_rate": 1e-3,
            "lambda": 1.0,
        }
    },
    alpha=1.0,
)

# context must be exactly context_dim=32 floats
arm, interaction_id = db.predict("content", context)
db.reward(interaction_id, reward)

Feature flag: NeuralLinUCB is compiled in only when the neural feature is enabled. The default binary contains no Candle or ML dependencies.

cargo build --release --features neural
cargo build --release --features neural,cuda
cargo build --release --features neural,metal

Hardware Acceleration

The compute device is selected at startup via the BANDITDB_DEVICE environment variable. When unset, BanditDB auto-detects: CUDA → Metal → CPU.

Value	Requires	Effect
unset or `auto`	—	CUDA if available, then Metal, then CPU.
`cuda`	`--features neural,cuda`	CUDA GPU 0.
`cuda:N`	`--features neural,cuda`	Specific CUDA device ordinal.
`metal`	`--features neural,metal`	Apple Silicon GPU.
`cpu`	`--features neural`	Force CPU regardless of available hardware.

Cold Start and Warm Restart

Retraining never blocks a prediction. The worker trains on a private copy of the network and swaps in the finished weights atomically when it is done; predictions read whichever snapshot is published and take no lock on the model at all.

This was a genuine outage in earlier versions, not a theoretical concern. A prediction held the arm lock and then wanted the neural lock, while a retrain held the neural lock and then wanted the arm lock. With a writer queued behind them, the task-fair lock closed the cycle and the server stopped serving — process alive, health check green, every request hung.

Cold start: a freshly created campaign uses randomly-initialised MLP weights until the first Algorithm 2 retrain fires at retrain_every rewards. Random projections are a usable baseline — LinUCB adapts even in a random embedding space — but expect lower sample efficiency than a trained network during this initial phase.

Warm restart: after each retrain, arm matrices (A⁻¹, b, θ) are re-accumulated by replaying the reward buffer through the updated network. This preserves most of the learned signal across retrains instead of resetting LinUCB statistics from scratch.

Recovery: MLP weights are saved as a .safetensors sidecar at {DATA_DIR}/neural/{campaign_id}.safetensors at each checkpoint and reloaded automatically on startup. If the sidecar is missing, the campaign recovers gracefully with fresh random weights and relearns from the WAL tail — model quality degrades temporarily but no data is lost.