olmo_tap.experiments.uncertainty.engine

Uncertainty finetuning protocol. See https://www.overleaf.com/read/kpnzybhdvwnh#a3aa13 for theory details.

Functions

get_calibration_prob(logits, config)
train(model, frozen_head_handler, ...[, ...])	Performs supervised uncertainty finetuning on a HydraTransformer model.

olmo_tap.experiments.uncertainty.engine.get_calibration_prob(logits: Tensor, config: TrainingConfig) → Tensor

olmo_tap.experiments.uncertainty.engine.train(model: HydraTransformer, frozen_head_handler: FrozenHeadHandler, exp_config: ExperimentConfig, optimizer: Optimizer, scheduler: LRScheduler, swap_every_n_steps: int = 100)

Performs supervised uncertainty finetuning on a HydraTransformer model. Assumed that 2 heads (uncertainty head at 0th index, frozen LLM head at 1st index) are loaded.

Parameters:

• model – HydraTransformer LLM model being finetuend.

• frozen_head_handler – Cycles through randomly selected frozen LLM heads to avoid uncertainty head overfitting to any one LLM head.

• optimizer – Any torch optim object.

• scheduler – Any torch scheduler object.

• swap_every_n_steps – How often to cycle through frozen LLM heads.

Parap exp_config:

Global config object storing experiment details.