diff --git a/examples/standard-model/symmetryIII_module.py b/examples/standard-model/symmetryIII_module.py new file mode 100644 index 00000000..74694f83 --- /dev/null +++ b/examples/standard-model/symmetryIII_module.py @@ -0,0 +1,295 @@ +import torch +import torch.nn.functional as F +from typing import Union, Optional +from sklearn.metrics import accuracy_score + +from torch_concepts.nn import MonotonicScoresEmbeddingToConcept, PrototypeConceptEmbeddingToConcept + + +class InstanceBasedArchitecture(torch.nn.Module): + """ + Complete instance-based architecture combining encoder and predictor. + + This is a convenience wrapper that combines MonotonicScoresEmbeddingToConcept + and PrototypeConceptEmbeddingToConcept into a single module. + + Args: + proto_samples: Tensor of shape [max_prototypes, num_concepts, n_features] - prototype feature vectors. + proto_scores: Tensor of shape [max_prototypes, num_concepts] - scores for sorting prototypes. + rank_dim: Dimension of the low-rank embedding for concepts. + temperature: Temperature for backward pass. + temp_forward: Temperature for forward pass. + use_straight_through: Use straight-through estimator. + learnable_prototypes: Whether prototypes should be learnable parameters. + + Example: + >>> proto_samples = torch.randn(10, 100, 50) + >>> proto_scores = torch.randn(10, 100) + >>> model = InstanceBasedArchitecture(proto_samples, proto_scores, rank_dim=32) + >>> x = torch.randn(32, 50) + >>> output = model(x) # [32, 100] + """ + def __init__( + self, + proto_samples: torch.Tensor, + proto_scores: torch.Tensor, + rank_dim: int = 32, + temperature: float = 1.0, + temp_forward: Optional[float] = None, + use_straight_through: bool = True, + learnable_prototypes: bool = False + ): + super().__init__() + + max_prototypes, num_concepts, n_features = proto_samples.shape + + self.embeddings = torch.nn.Embedding(max_prototypes, rank_dim) + + self.encoder = MonotonicScoresEmbeddingToConcept( + in_embeddings=rank_dim, + out_concepts=num_concepts, + ) + + self.predictor = PrototypeConceptEmbeddingToConcept( + proto_samples=proto_samples, + proto_scores=proto_scores, + out_concepts=num_concepts, + learnable_prototypes=learnable_prototypes, + temperature=temperature, + temp_forward=temp_forward, + use_straight_through=use_straight_through + ) + + self.num_concepts = num_concepts + self.max_prototypes = max_prototypes + self.n_features = n_features + + # Expose prototypes for compatibility + self.prototypes = self.predictor.prototypes + self.embedding = self.embeddings + self.projection = self.encoder.projection + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass. + + Args: + x: Tensor of shape [batch, n_features] - input features. + + Returns: + torch.Tensor: Tensor of shape [batch, num_concepts] - concept predictions. + """ + # Generate concept weights + concepts = self.encoder(self.embeddings.weight) # [1, num_concepts, max_prototypes] + + # Aggregate with embeddings + output = self.predictor(concepts, x) # [batch, num_concepts] + + return output + + +def main(): + import time + torch.manual_seed(42) + + print("="*70) + print(f"Train") + print("="*70) + print() + n_features = 50 + batch_size = 512 + rank_dim = 10 + max_prototypes = 10 + num_concepts = 1000 + proto_samples = torch.randn(max_prototypes, num_concepts, n_features) + proto_scores = torch.randn(max_prototypes, num_concepts) + x_train = torch.randn(batch_size, n_features) + y_train = ((x_train[:, 2] + x_train[:, 3])>0).float().unsqueeze(1) + + cum_mlp = InstanceBasedArchitecture(proto_samples, proto_scores, rank_dim=rank_dim) + predictor = torch.nn.Linear(num_concepts, 1) + model = torch.nn.Sequential(cum_mlp, predictor) + + # Forward + backward + optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + for epoch in range(100): + optimizer.zero_grad() + output = model(x_train) + loss = F.binary_cross_entropy_with_logits(output, y_train) + loss.backward() + optimizer.step() + if epoch % 10 == 0: + task_accuracy = accuracy_score(y_train, output.detach().cpu().numpy()>0) + print(f"Epoch: {epoch}, loss: {loss}, acc.: {task_accuracy}") + + print() + print("="*70) + print("InstanceBasedArchitecture Performance Benchmark") + print("="*70) + print() + + # Test configurations + n_features = 10 + batch_size = 32 + rank_dim = 32 + prototype_configs = [10, 100] + concept_configs = [10, 100, 1000, 10000] + + print(f"Configuration:") + print(f" Batch size: {batch_size}") + print(f" Feature dimension: {n_features}") + print(f" Rank dimension: {rank_dim}") + print() + + # Results table + print(f"{'Prototypes':<12} {'Concepts':<10} {'Forward (ms)':<15} {'Memory (MB)':<15} {'Throughput':<15}") + print("-" * 70) + + for max_prototypes in prototype_configs: + for num_concepts in concept_configs: + # Create prototype samples and scores + proto_samples = torch.randn(max_prototypes, num_concepts, n_features) + proto_scores = torch.randn(max_prototypes, num_concepts) + + # Create model + model = InstanceBasedArchitecture(proto_samples, proto_scores, rank_dim=rank_dim) + model.eval() + + # Create test input + x = torch.randn(batch_size, n_features) + + # Warmup + with torch.no_grad(): + for _ in range(5): + _ = model(x) + + # Benchmark forward pass + num_runs = 10 + + torch.cuda.synchronize() if torch.cuda.is_available() else None + start_time = time.time() + + with torch.no_grad(): + for _ in range(num_runs): + output = model(x) + + torch.cuda.synchronize() if torch.cuda.is_available() else None + end_time = time.time() + + # Calculate metrics + avg_time_ms = (end_time - start_time) / num_runs * 1000 + + # Memory estimate (rough) + param_memory = sum(p.numel() * 4 for p in model.parameters()) / (1024 * 1024) # MB + proto_memory = model.prototypes.numel() * 4 / (1024 * 1024) # MB + total_memory = param_memory + proto_memory + + # Throughput (samples per second) + throughput = batch_size * num_runs / (end_time - start_time) + + print(f"{max_prototypes:<12} {num_concepts:<10} {avg_time_ms:<15.3f} " + f"{total_memory:<15.2f} {throughput:<15.0f}") + + print() + print("="*70) + print("Detailed Test: 10 prototypes, 1k concepts") + print("="*70) + + max_prototypes = 10 + num_concepts = 1000 + + proto_samples = torch.randn(max_prototypes, num_concepts, n_features) + proto_scores = torch.randn(max_prototypes, num_concepts) + + model = InstanceBasedArchitecture(proto_samples, proto_scores, rank_dim=rank_dim) + + print(f"\nModel Statistics:") + print(f" Total parameters: {sum(p.numel() for p in model.parameters()):,}") + print(f" - Embedding: {model.embedding.weight.numel():,}") + print(f" - Projection weight: {model.projection.weight.numel():,}") + print(f" - Projection bias: {model.projection.bias.numel():,}") + print(f" Prototype storage: {model.prototypes.numel():,} elements") + print(f" Prototype memory: {model.prototypes.numel() * 4 / (1024**2):.2f} MB") + print() + + # Test with different batch sizes + print("Batch Size Scaling:") + print(f"{'Batch Size':<12} {'Forward (ms)':<15} {'Per Sample (ms)':<15}") + print("-" * 45) + + for bs in [1, 8, 32, 64]: + x_test = torch.randn(bs, n_features) + + # Warmup + with torch.no_grad(): + for _ in range(3): + _ = model(x_test) + + # Benchmark + num_runs = 5 + start_time = time.time() + with torch.no_grad(): + for _ in range(num_runs): + _ = model(x_test) + end_time = time.time() + + avg_time_ms = (end_time - start_time) / num_runs * 1000 + per_sample_ms = avg_time_ms / bs + + print(f"{bs:<12} {avg_time_ms:<15.3f} {per_sample_ms:<15.4f}") + + print() + print("="*70) + print("Gradient Computation Test") + print("="*70) + + # Test gradient computation overhead + x_train = torch.randn(batch_size, n_features) + y_train = torch.randn(batch_size, num_concepts) + + # Forward only + start_time = time.time() + for _ in range(10): + with torch.no_grad(): + output = model(x_train) + forward_time = (time.time() - start_time) / 10 * 1000 + + # Forward + backward + optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + start_time = time.time() + for _ in range(10): + optimizer.zero_grad() + output = model(x_train) + loss = F.mse_loss(output, y_train) + loss.backward() + optimizer.step() + train_time = (time.time() - start_time) / 10 * 1000 + + print(f"\nWith {num_concepts:,} concepts, {max_prototypes} prototypes:") + print(f" Forward pass: {forward_time:.2f} ms") + print(f" Forward + Backward: {train_time:.2f} ms") + print(f" Backward overhead: {train_time - forward_time:.2f} ms ({(train_time/forward_time - 1)*100:.1f}% increase)") + + print() + print("="*70) + print("Monotonicity Verification (1k concepts sample)") + print("="*70) + + # Check monotonicity for a few concepts + with torch.no_grad(): + for c in range(min(3, num_concepts)): + proto_inputs = model.prototypes[c] # [max_prototypes, n_features] + proto_outputs = model(proto_inputs)[:, c] # Predictions for concept c + + diffs = torch.diff(proto_outputs) + is_monotonic = torch.all(diffs >= -1e-6).item() + + print(f"\nConcept {c}:") + print(f" Monotonic: {is_monotonic}") + print(f" Min diff: {diffs.min().item():.6f}") + print(f" Max diff: {diffs.max().item():.6f}") + print(f" Mean diff: {diffs.mean().item():.6f}") + + +if __name__ == "__main__": + main() diff --git a/torch_concepts/nn/__init__.py b/torch_concepts/nn/__init__.py index b14e1267..99bfd8dd 100644 --- a/torch_concepts/nn/__init__.py +++ b/torch_concepts/nn/__init__.py @@ -25,12 +25,17 @@ # Encoders from .modules.low.encoders.linear import LinearEmbeddingToConcept +from .modules.low.encoders.prototype import MonotonicScoresEmbeddingToConcept # Predictors from .modules.low.predictors.call import CallableConceptToConcept from .modules.low.predictors.hypernet import HyperlinearConceptEmbeddingToConcept from .modules.low.predictors.linear import LinearConceptToConcept from .modules.low.predictors.mix import MixConceptEmbeddingToConcept +from .modules.low.predictors.prototype import PrototypeConceptEmbeddingToConcept + +# Ops +from .modules.low.ops import StraightThroughSoftmax # Dense layers from .modules.low.dense_layers import Dense, ResidualMLP, MLP, LinearEmbeddingEncoder, SelectorEmbeddingEncoder @@ -118,12 +123,14 @@ # Encoder classes "LinearEmbeddingToConcept", + "MonotonicScoresEmbeddingToConcept", # Predictor classes "LinearConceptToConcept", "CallableConceptToConcept", "HyperlinearConceptEmbeddingToConcept", "MixConceptEmbeddingToConcept", + "PrototypeConceptEmbeddingToConcept", # Dense layers "Dense", diff --git a/torch_concepts/nn/modules/low/encoders/prototype.py b/torch_concepts/nn/modules/low/encoders/prototype.py new file mode 100644 index 00000000..bffc2be3 --- /dev/null +++ b/torch_concepts/nn/modules/low/encoders/prototype.py @@ -0,0 +1,59 @@ +import torch +import torch.nn.functional as F +from typing import Union + +from .....annotations import Annotations +from ..base.layer import BaseConceptLayer + + +class MonotonicScoresEmbeddingToConcept(BaseConceptLayer): + """ + Generates monotonic concept scores. + + This encoder produces concept scores where each sample is represented by a cumulative sum of positive scores, + ensuring monotonicity across samples. + + Args: + in_embeddings: Dimension of the input embeddings (in_embeddings). + out_concepts: Number of concepts to generate (num_concepts). + + Example: + >>> embeddings = torch.randn(32, 10) # Batch of 32 embeddings with 10 features + >>> encoder = MonotonicScoresEmbeddingToConcept(out_concepts=100, in_embeddings=10) + >>> concepts = encoder(embeddings) # [32, 100] + """ + def __init__( + self, + in_embeddings: Union[int, Annotations], + out_concepts: Union[int, Annotations], + *args, + **kwargs + ): + super().__init__( + out_concepts=out_concepts, + in_concepts=None, + in_embeddings=in_embeddings + ) + self.projection = torch.nn.Linear( + self.in_embeddings_shape, + self.out_concepts_shape + ) + + def forward( + self, + embeddings: torch.Tensor + ) -> torch.Tensor: + """ + Generate cumulative weights for all concepts. + + Returns: + torch.Tensor: Tensor of shape [max_prototypes, num_concepts] + """ + scores = self.projection(embeddings) # [max_prototypes, num_concepts] + + # Ensure positivity with softplus (numerically stable) + positive_scores = F.softplus(scores) + # Cumulative sum ensures monotonicity + cumulative_output = torch.cumsum(positive_scores, dim=0) # [max_prototypes, num_concepts] + + return cumulative_output diff --git a/torch_concepts/nn/modules/low/ops.py b/torch_concepts/nn/modules/low/ops.py index 19913f36..1bdbf6b5 100644 --- a/torch_concepts/nn/modules/low/ops.py +++ b/torch_concepts/nn/modules/low/ops.py @@ -1,5 +1,37 @@ import torch from torch import nn +import torch.nn.functional as F + + +class StraightThroughSoftmax(torch.autograd.Function): + """ + Straight-Through Estimator for softmax: + Forward pass uses sharp temperature (peaked selection) + Backward pass uses soft temperature (informative gradients) + """ + @staticmethod + def forward(ctx, logits, temp_forward, temp_backward, dim): + ctx.temp_backward = temp_backward + ctx.dim = dim + ctx.save_for_backward(logits) + # Forward: very peaked distribution + return F.softmax(logits / temp_forward, dim=dim) + + @staticmethod + def backward(ctx, grad_output): + logits, = ctx.saved_tensors + temp = ctx.temp_backward + dim = ctx.dim + + # Backward: use softer temperature for better gradient flow + probs = F.softmax(logits / temp, dim=dim) + + # Compute softmax Jacobian-vector product + grad_input = probs * (grad_output - (grad_output * probs).sum(dim=dim, keepdim=True)) + grad_input = grad_input / temp + + return grad_input, None, None, None + class SumOp(nn.Module): r"""Sum ``n_terms`` equal-size contributions concatenated along the last dim. diff --git a/torch_concepts/nn/modules/low/predictors/prototype.py b/torch_concepts/nn/modules/low/predictors/prototype.py new file mode 100644 index 00000000..5d34051a --- /dev/null +++ b/torch_concepts/nn/modules/low/predictors/prototype.py @@ -0,0 +1,161 @@ +import torch +from typing import Optional, Union +import torch.nn.functional as F + +from .....annotations import Annotations +from ..base .layer import BaseConceptLayer +from ..ops import StraightThroughSoftmax + + +class PrototypeConceptEmbeddingToConcept(BaseConceptLayer): + """ + Aggregates prototype-based concept weights with input embeddings. + + This predictor computes similarity between input embeddings and stored prototypes, + then uses concept weights from the encoder to produce final concept predictions. + + Args: + out_concepts: Number of output concepts (num_concepts). + proto_samples: Tensor of shape [max_prototypes, num_concepts, n_features] - prototype feature vectors. + proto_scores: Tensor of shape [max_prototypes, num_concepts] - scores for sorting prototypes. + learnable_prototypes: Whether prototypes should be learnable parameters. + temperature: Temperature for backward pass (default: 1.0). + temp_forward: Temperature for forward pass (default: 0.01). + use_straight_through: Use straight-through estimator for peaked forward, soft backward. + + Example: + >>> proto_samples = torch.randn(10, 100, 50) # 10 prototypes, 100 concepts, 50 features + >>> proto_scores = torch.randn(10, 100) + >>> predictor = PrototypeConceptEmbeddingToConcept( + ... out_concepts=100, + ... proto_samples=proto_samples, + ... proto_scores=proto_scores, + ... ) + >>> concepts = torch.randn(10, 100) # From encoder + >>> embeddings = torch.randn(32, 50) # Batch of 32 + >>> output = predictor(concepts, embeddings) # [32, 100] + """ + def __init__( + self, + out_concepts: Union[int, Annotations], + proto_samples: torch.Tensor, + proto_scores: torch.Tensor, + learnable_prototypes: bool = False, + temperature: float = 1.0, + temp_forward: Optional[float] = None, + use_straight_through: bool = True, + *args, + **kwargs + ): + max_prototypes, num_concepts, n_features = proto_samples.shape + # Infer in_embeddings + in_embeddings = n_features + + super().__init__( + out_concepts=out_concepts, + in_concepts=None, # Concepts come from encoder, not traditional input + in_embeddings=in_embeddings + ) + + assert self.out_concepts_shape == num_concepts, \ + f"out_concepts ({self.out_concepts_shape}) must match num_concepts from proto_samples ({num_concepts})" + + self.num_concepts = num_concepts + self.max_prototypes = max_prototypes + self.n_features = n_features + + # Temperature settings + self.use_straight_through = use_straight_through + if use_straight_through: + self.temp_forward = temp_forward if temp_forward is not None else 0.01 + self.temp_backward = temperature + else: + self.temp_forward = temperature + self.temp_backward = temperature + + self.register_buffer('temperature_forward', torch.tensor(self.temp_forward)) + self.register_buffer('temperature_backward', torch.tensor(self.temp_backward)) + + # Sort prototypes by scores (ascending order for monotonicity) + sorted_indices = torch.stack([ + torch.argsort(proto_scores[:, i], descending=False) + for i in range(num_concepts) + ]) + + # Reorder prototypes: [num_concepts, max_prototypes, n_features] + prototypes = torch.stack([ + proto_samples[sorted_indices[i], i, :] + for i in range(num_concepts) + ]) + + # Store or register as parameter + if learnable_prototypes: + self.prototypes = torch.nn.Parameter(prototypes) + else: + self.register_buffer('prototypes', prototypes) + + def forward( + self, + concepts: torch.Tensor, + embeddings: torch.Tensor + ) -> torch.Tensor: + """ + Aggregate concept weights with embeddings via prototype similarity. + + Args: + concepts: Tensor of shape [max_prototypes, num_concepts]. + embeddings: Tensor of shape [batch, n_features]. + + Returns: + torch.Tensor: Tensor of shape [batch, num_concepts] - final concept predictions. + """ + similarity = self.similarity_scores(concepts, embeddings) + + # Weighted sum over prototypes + # similarity: [batch, num_concepts, max_prototypes] + # cumulative_weights: [num_concepts, max_prototypes] + output = (similarity * concepts.T.unsqueeze(0)).sum(dim=2) # [batch, num_concepts] + + return output + + def similarity_scores( + self, + concepts: torch.Tensor, + embeddings: torch.Tensor + ) -> torch.Tensor: + """ + Compute similarity scores between embeddings and prototypes. + """ + + assert self.max_prototypes == concepts.shape[0], \ + f"Expected concepts to have shape [{self.max_prototypes}, {self.num_concepts}], got {concepts.shape}" + + batch_size = embeddings.shape[0] + + # Compute similarities between embeddings and prototypes + # embeddings: [batch, n_features] -> [batch, 1, 1, n_features] + # prototypes: [num_concepts, max_prototypes, n_features] -> [1, num_concepts, max_prototypes, n_features] + x_expanded = embeddings.view(batch_size, 1, 1, self.n_features) + proto_expanded = self.prototypes.unsqueeze(0) + + # Compute squared distances: [batch, num_concepts, max_prototypes] + squared_distances = torch.sum((x_expanded - proto_expanded) ** 2, dim=-1) + + # Use negative squared distances as logits + logits = -squared_distances + + # Compute similarities with appropriate method + if self.use_straight_through: + # Straight-through: peaked forward, smooth backward + similarity = StraightThroughSoftmax.apply( + logits, + self.temperature_forward, + self.temperature_backward, + 2 # dim for softmax + ) + else: + # Standard softmax with single temperature + temp = torch.clamp(self.temperature_forward, min=0.1) + similarity = F.softmax(logits / temp, dim=2) + + return similarity