osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
406 lines
14 KiB
Rust
406 lines
14 KiB
Rust
//! Cross-Entropy Loss Implementation
|
||
//!
|
||
//! Provides efficient cross-entropy loss computation for language modeling tasks.
|
||
//! Supports both standard cross-entropy and sparse cross-entropy with label smoothing.
|
||
|
||
use crate::losses::{Loss, Reduction};
|
||
use crate::{Result, TransformerError};
|
||
use rtx_tensor::Tensor;
|
||
use std::f32::NEG_INFINITY;
|
||
use tracing::{debug, trace};
|
||
|
||
/// Cross-entropy loss configuration
|
||
#[derive(Debug, Clone)]
|
||
pub struct CrossEntropyConfig {
|
||
/// Reduction method for loss computation
|
||
pub reduction: Reduction,
|
||
/// Label smoothing factor (0.0 = no smoothing, 0.1 = 10% smoothing)
|
||
pub label_smoothing: f32,
|
||
/// Whether to ignore index -100 (padding tokens)
|
||
pub ignore_index: Option<i64>,
|
||
}
|
||
|
||
impl Default for CrossEntropyConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
reduction: Reduction::Mean,
|
||
label_smoothing: 0.0,
|
||
ignore_index: Some(-100),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Cross-entropy loss implementation for language modeling
|
||
///
|
||
/// Computes the cross-entropy loss between predictions and targets:
|
||
/// loss = -sum(target * `log_softmax(predictions)`)
|
||
///
|
||
/// For language modeling, typically:
|
||
/// - predictions: [`batch_size`, `seq_len`, `vocab_size`] logits
|
||
/// - targets: [`batch_size`, `seq_len`] token indices
|
||
pub struct CrossEntropyLoss {
|
||
config: CrossEntropyConfig,
|
||
}
|
||
|
||
impl CrossEntropyLoss {
|
||
/// Create a new cross-entropy loss with default configuration
|
||
#[must_use]
|
||
pub fn new() -> Self {
|
||
Self {
|
||
config: CrossEntropyConfig::default(),
|
||
}
|
||
}
|
||
|
||
/// Create a new cross-entropy loss with custom configuration
|
||
#[must_use]
|
||
pub fn with_config(config: CrossEntropyConfig) -> Self {
|
||
Self { config }
|
||
}
|
||
|
||
/// Get the configuration
|
||
#[must_use]
|
||
pub fn config(&self) -> &CrossEntropyConfig {
|
||
&self.config
|
||
}
|
||
|
||
/// Compute log-softmax of logits
|
||
///
|
||
/// For numerical stability, we use the log-sum-exp trick:
|
||
/// `log_softmax(x)` = x - log(sum(exp(x - max(x))))
|
||
fn log_softmax(&self, logits: &Tensor) -> Result<Tensor> {
|
||
let dims = logits.shape().dims();
|
||
let vocab_dim = dims.len() - 1; // Last dimension is vocab_size
|
||
|
||
// Find maximum along vocab dimension for numerical stability
|
||
let max_logits = self.max_along_dim(logits, vocab_dim)?;
|
||
|
||
// Subtract max for numerical stability: x - max(x)
|
||
let shifted_logits = self.subtract_broadcasted(logits, &max_logits)?;
|
||
|
||
// Compute exp(x - max(x))
|
||
let exp_shifted = self.exp_tensor(&shifted_logits)?;
|
||
|
||
// Sum along vocab dimension: sum(exp(x - max(x)))
|
||
let sum_exp = self.sum_along_dim(&exp_shifted, vocab_dim)?;
|
||
|
||
// Compute log(sum(exp(x - max(x))))
|
||
let log_sum_exp = self.log_tensor(&sum_exp)?;
|
||
|
||
// Return log_softmax: (x - max(x)) - log(sum(exp(x - max(x))))
|
||
self.subtract_broadcasted(&shifted_logits, &log_sum_exp)
|
||
}
|
||
|
||
/// Compute sparse cross-entropy loss
|
||
///
|
||
/// For each position (b, t), we compute:
|
||
/// loss[b, t] = -`log_softmax`[b, t, target[b, t]]
|
||
fn sparse_cross_entropy(&self, log_probs: &Tensor, targets: &Tensor) -> Result<Tensor> {
|
||
let batch_size = log_probs.shape()[0];
|
||
let seq_len = log_probs.shape()[1];
|
||
let vocab_size = log_probs.shape()[2];
|
||
|
||
// Convert to CPU for indexing operations
|
||
let log_probs_cpu = log_probs.to_cpu()?;
|
||
let targets_cpu = targets.to_cpu()?;
|
||
|
||
let mut losses = Vec::with_capacity(batch_size * seq_len);
|
||
|
||
for b in 0..batch_size {
|
||
for t in 0..seq_len {
|
||
// Manual indexing into flat vector: [b, t] -> b * seq_len + t
|
||
let target_idx = targets_cpu[b * seq_len + t] as i64;
|
||
|
||
let loss_value = if let Some(ignore_idx) = self.config.ignore_index {
|
||
if target_idx == ignore_idx {
|
||
0.0 // Ignore this position
|
||
} else {
|
||
// Manual indexing: [b, t, c] -> (b * seq_len + t) * vocab_size + c
|
||
let log_prob =
|
||
log_probs_cpu[(b * seq_len + t) * vocab_size + target_idx as usize];
|
||
-log_prob
|
||
}
|
||
} else {
|
||
// Manual indexing: [b, t, c] -> (b * seq_len + t) * vocab_size + c
|
||
let log_prob =
|
||
log_probs_cpu[(b * seq_len + t) * vocab_size + target_idx as usize];
|
||
-log_prob
|
||
};
|
||
|
||
losses.push(loss_value);
|
||
}
|
||
}
|
||
|
||
// Create tensor with per-token losses
|
||
let device = log_probs.device();
|
||
Tensor::from_data(losses, vec![batch_size, seq_len], device)
|
||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||
}
|
||
|
||
/// Apply label smoothing
|
||
///
|
||
/// Smoothed loss = (1 - α) * `ce_loss` + α * `uniform_loss`
|
||
/// where α is the smoothing factor
|
||
fn apply_label_smoothing(&self, ce_loss: &Tensor, log_probs: &Tensor) -> Result<Tensor> {
|
||
if self.config.label_smoothing == 0.0 {
|
||
return Ok(ce_loss.clone());
|
||
}
|
||
|
||
let vocab_size = log_probs.shape()[2] as f32;
|
||
let smoothing = self.config.label_smoothing;
|
||
|
||
// Uniform distribution loss: -mean(log_probs)
|
||
let uniform_loss = self.mean_along_dim(log_probs, 2)?;
|
||
let uniform_loss = self.negate_tensor(&uniform_loss)?;
|
||
|
||
// Combine: (1 - smoothing) * ce_loss + smoothing * uniform_loss
|
||
let ce_weight = 1.0 - smoothing;
|
||
let uniform_weight = smoothing / vocab_size;
|
||
|
||
let weighted_ce = self.scale_tensor(ce_loss, ce_weight)?;
|
||
let weighted_uniform = self.scale_tensor(&uniform_loss, uniform_weight)?;
|
||
|
||
self.add_tensors(&weighted_ce, &weighted_uniform)
|
||
}
|
||
|
||
/// Apply reduction to loss tensor
|
||
fn apply_reduction(&self, loss: &Tensor) -> Result<Tensor> {
|
||
match self.config.reduction {
|
||
Reduction::None => Ok(loss.clone()),
|
||
Reduction::Mean => {
|
||
let total_elements = loss.numel() as f32;
|
||
let sum_loss = self.sum_all(loss)?;
|
||
self.scale_tensor(&sum_loss, 1.0 / total_elements)
|
||
}
|
||
Reduction::Sum => self.sum_all(loss),
|
||
}
|
||
}
|
||
|
||
// Helper methods for tensor operations
|
||
// These would ideally be implemented as proper tensor operations
|
||
|
||
fn max_along_dim(&self, tensor: &Tensor, dim: usize) -> Result<Tensor> {
|
||
// Simplified implementation - in practice, use proper tensor max operation
|
||
let data = tensor.to_cpu()?;
|
||
let shape = tensor.shape().dims();
|
||
|
||
if dim >= shape.len() {
|
||
return Err(TransformerError::tensor_op(
|
||
"Invalid dimension for max operation".to_string(),
|
||
));
|
||
}
|
||
|
||
let reduced_shape: Vec<usize> = shape
|
||
.iter()
|
||
.enumerate()
|
||
.filter_map(|(i, &size)| if i == dim { None } else { Some(size) })
|
||
.collect();
|
||
|
||
let max_val = data.iter().fold(NEG_INFINITY, |a, &b| a.max(b));
|
||
let result_size = reduced_shape.iter().product();
|
||
let result_data = vec![max_val; result_size];
|
||
|
||
Tensor::from_data(result_data, reduced_shape, tensor.device())
|
||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||
}
|
||
|
||
fn subtract_broadcasted(&self, a: &Tensor, b: &Tensor) -> Result<Tensor> {
|
||
// Simplified - in practice, use proper broadcasting
|
||
(a.clone() - b.clone()).map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||
}
|
||
|
||
fn exp_tensor(&self, tensor: &Tensor) -> Result<Tensor> {
|
||
tensor
|
||
.exp()
|
||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||
}
|
||
|
||
fn log_tensor(&self, tensor: &Tensor) -> Result<Tensor> {
|
||
tensor
|
||
.log()
|
||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||
}
|
||
|
||
fn sum_along_dim(&self, tensor: &Tensor, dim: usize) -> Result<Tensor> {
|
||
tensor
|
||
.sum(Some(dim))
|
||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||
}
|
||
|
||
fn mean_along_dim(&self, tensor: &Tensor, dim: usize) -> Result<Tensor> {
|
||
let sum = self.sum_along_dim(tensor, dim)?;
|
||
let size = tensor.shape()[dim] as f32;
|
||
self.scale_tensor(&sum, 1.0 / size)
|
||
}
|
||
|
||
fn negate_tensor(&self, tensor: &Tensor) -> Result<Tensor> {
|
||
tensor
|
||
.mul_scalar(-1.0)
|
||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||
}
|
||
|
||
fn scale_tensor(&self, tensor: &Tensor, scale: f32) -> Result<Tensor> {
|
||
tensor
|
||
.mul_scalar(scale)
|
||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||
}
|
||
|
||
fn add_tensors(&self, a: &Tensor, b: &Tensor) -> Result<Tensor> {
|
||
a.add(b)
|
||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||
}
|
||
|
||
fn sum_all(&self, tensor: &Tensor) -> Result<Tensor> {
|
||
tensor
|
||
.sum(None)
|
||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||
}
|
||
}
|
||
|
||
impl Loss for CrossEntropyLoss {
|
||
fn forward(&self, predictions: &Tensor, targets: &Tensor) -> Result<Tensor> {
|
||
trace!(
|
||
"Computing cross-entropy loss: predictions={:?}, targets={:?}",
|
||
predictions.shape(),
|
||
targets.shape()
|
||
);
|
||
|
||
// Validate input shapes
|
||
let pred_shape = predictions.shape().dims();
|
||
let target_shape = targets.shape().dims();
|
||
|
||
if pred_shape.len() != 3 {
|
||
return Err(TransformerError::tensor_op(format!(
|
||
"Predictions must be 3D [batch, seq, vocab], got shape {pred_shape:?}"
|
||
)));
|
||
}
|
||
|
||
if target_shape.len() != 2 {
|
||
return Err(TransformerError::tensor_op(format!(
|
||
"Targets must be 2D [batch, seq], got shape {target_shape:?}"
|
||
)));
|
||
}
|
||
|
||
if pred_shape[0] != target_shape[0] || pred_shape[1] != target_shape[1] {
|
||
return Err(TransformerError::tensor_op(format!(
|
||
"Batch/sequence dimensions mismatch: pred {:?} vs target {:?}",
|
||
&pred_shape[..2],
|
||
target_shape
|
||
)));
|
||
}
|
||
|
||
// Compute log-softmax
|
||
let log_probs = self.log_softmax(predictions)?;
|
||
|
||
// Compute sparse cross-entropy loss
|
||
let mut ce_loss = self.sparse_cross_entropy(&log_probs, targets)?;
|
||
|
||
// Apply label smoothing if configured
|
||
if self.config.label_smoothing > 0.0 {
|
||
ce_loss = self.apply_label_smoothing(&ce_loss, &log_probs)?;
|
||
}
|
||
|
||
// Apply reduction
|
||
let final_loss = self.apply_reduction(&ce_loss)?;
|
||
|
||
debug!(
|
||
"Cross-entropy loss computed: reduction={:?}, value={:.4}",
|
||
self.config.reduction,
|
||
final_loss
|
||
.to_cpu()
|
||
.unwrap_or_else(|_| vec![0.0])
|
||
.first()
|
||
.copied()
|
||
.unwrap_or(0.0)
|
||
);
|
||
|
||
Ok(final_loss)
|
||
}
|
||
|
||
fn reduction(&self) -> Reduction {
|
||
self.config.reduction
|
||
}
|
||
|
||
fn name(&self) -> &'static str {
|
||
"CrossEntropyLoss"
|
||
}
|
||
}
|
||
|
||
#[cfg(all(test, feature = "disabled_tests"))]
|
||
mod tests {
|
||
use super::*;
|
||
use rtx_tensor::Device;
|
||
|
||
#[test]
|
||
fn test_cross_entropy_basic() {
|
||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||
let loss_fn = CrossEntropyLoss::new();
|
||
|
||
// Create simple test case
|
||
let predictions = Tensor::from_data(
|
||
vec![
|
||
// Batch 1, Token 1: high prob for class 0
|
||
1.0, 0.0, 0.0, // Batch 1, Token 2: high prob for class 1
|
||
0.0, 1.0, 0.0,
|
||
],
|
||
vec![1, 2, 3], // [batch=1, seq=2, vocab=3]
|
||
&device,
|
||
)
|
||
.unwrap();
|
||
|
||
let targets = Tensor::from_data(
|
||
vec![0i64, 1i64], // Correct predictions
|
||
vec![1, 2], // [batch=1, seq=2]
|
||
&device,
|
||
)
|
||
.unwrap();
|
||
|
||
let loss = loss_fn.forward(&predictions, &targets).unwrap();
|
||
let loss_value: f32 = loss.to_cpu().unwrap()[0];
|
||
|
||
// Loss should be low for correct predictions
|
||
assert!(loss_value >= 0.0, "Loss should be non-negative");
|
||
assert!(
|
||
loss_value < 2.0,
|
||
"Loss should be reasonable for correct predictions"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_cross_entropy_shapes() {
|
||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||
let loss_fn = CrossEntropyLoss::new();
|
||
|
||
// Wrong prediction shape (2D instead of 3D)
|
||
let predictions = Tensor::zeros(&[2, 3], &device).unwrap();
|
||
let targets = Tensor::zeros_typed(&[2], DType::I64, &device).unwrap();
|
||
|
||
let result = loss_fn.forward(&predictions, &targets);
|
||
assert!(result.is_err(), "Should fail with wrong prediction shape");
|
||
|
||
// Wrong target shape (3D instead of 2D)
|
||
let predictions = Tensor::zeros(&[2, 3, 4], &device).unwrap();
|
||
let targets = Tensor::zeros_typed(&[2, 3, 1], DType::I64, &device).unwrap();
|
||
|
||
let result = loss_fn.forward(&predictions, &targets);
|
||
assert!(result.is_err(), "Should fail with wrong target shape");
|
||
}
|
||
|
||
#[test]
|
||
fn test_label_smoothing() {
|
||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||
let config = CrossEntropyConfig {
|
||
label_smoothing: 0.1,
|
||
..Default::default()
|
||
};
|
||
let loss_fn = CrossEntropyLoss::with_config(config);
|
||
|
||
let predictions =
|
||
Tensor::from_data(vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0], vec![1, 2, 3], &device).unwrap();
|
||
|
||
let targets = Tensor::from_data(vec![0i64, 1i64], vec![1, 2], &device).unwrap();
|
||
|
||
let loss = loss_fn.forward(&predictions, &targets);
|
||
assert!(loss.is_ok(), "Label smoothing should work");
|
||
}
|
||
}
|