Softmax turns a model’s raw scores into probabilities: every value becomes nonnegative, all values add to 1, and the result can be read as the likelihood of each possible label or token. For classification, the highest probability is typically the prediction. section 5.2.1
Apply softmax after inference, once output bytes have been converted to 32-bit floating-point values. Exponentiate each score, add the exponentials, and divide every exponential by that total. This amplifies differences between predicted scores while normalizing them into a probability distribution. section 5.2.2
let exp_output = output_array.mapv(f32::exp);
let sum_exp = exp_output.sum();
let softmax_output = exp_output / sum_exp;
This implementation maps f32::exp over model outputs, sums the results, and divides the vector by that sum. section 5.2.2
Use the resulting probabilities according to the task:
- Image classification: pair probabilities with class indexes, sort them in descending order, and display the top predictions with their labels. section 5.2.2
- Token generation: apply softmax to the final token’s logits, then pass the probabilities to WeightedIndex to sample a likely next token with randomness. section 5.4.2
Reusable function for next-token selection
A story generator isolates the logits for the last token, sends that slice to a reusable softmax function, and samples one next-token index from the returned probability vector. The selected token is appended to the growing input sequence before decoding it back to text. section 5.4.2 section 5.4.2
fn softmax(logits: &[f32]) -> Vec<f32> {
let logits_array = Array::from_vec(logits.to_vec());
let exp_logits = logits_array.mapv(f32::exp);
let sum_exp = exp_logits.sum();
(exp_logits / sum_exp).to_vec()
}
The standalone Rust function accepts a logits slice and returns normalized probabilities, making the same transformation reusable during token generation. section 5.4.2
