RLAIF › Key Parameters

Understanding the key parameters of our implementation helps you adapt it to production needs. These are not arbitrary choices but design decisions that involve real tradeoffs.

The key parameters for the reward model and training loop are:

• vocab_size: The size of the vocabulary for the embedding layer. This determines how many unique tokens the model can stand for. In practice, this matches the tokenizer's vocabulary size. Using a real subword tokenizer (such as the BPE tokenizer described in the tokenization chapters) rather than the hash-based approach here would make vocab_size a well-defined quantity derived from the tokenizer's configuration.
• embed_dim: The dimensionality of the word embeddings. Higher values allow richer representations but increase computational cost. Typical values range from 64 to 768 for production reward models; we use 64 here to keep the demonstration fast.
• hidden_dim: The number of features in the hidden state of the LSTM. This controls the capacity of the encoder to capture sequential patterns. Since we use a bidirectional LSTM, the final representation has dimension $2 \times \text{hidden\_dim}$.
• epochs: The number of full passes through the training dataset. More epochs allow for better convergence but risk overfitting, especially with small datasets. With AI-generated data, you can generate more data to avoid overfitting rather than running more epochs.
• lr: The learning rate for the Adam optimizer. This controls how quickly the model updates its parameters. Values between $10^{-4}$ and $10^{-3}$ are common starting points for reward model training.
• batch_size: The number of training samples processed before updating the model parameters. Larger batches give more stable gradients but require more memory. For reward models trained on AI preferences, batch sizes of 32-128 are common in production.
• min_confidence: The confidence threshold below which preference labels are discarded. Setting this too high wastes data by rejecting borderline cases; setting it too low poisons the training signal with noisy labels. Values between 0.6 and 0.8 are typical.