Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,405 @@
|
||||
//! 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");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user