DOCSEARCHbyAlgolia
Submit Docs

go-attention

repository·main·Indexed 19 days ago

https://github.com/takara-ai/go-attention

A high-performance, pure Go implementation of attention mechanisms and transformer layers. It provides Dot-Product Attention, Multi-Head Attention (MHA), and full Transformer Layers including layer normalization and feed-forward networks. Designed for production reliability and edge computing, the library includes features for performance monitoring, memory pooling for Vector and Matrix types, and parallel processing to reduce GC pressure and maximize throughput without external dependencies.

Tokens
6.2K
Snippets
24
Records
27
Agent score
64%

What's inside go-attention

  1. Core Types in go-attention

    main

    The library uses three primary types for numerical data:

    • Vector: A 1D slice of float64 values used for embeddings and attention weights.
    • Matrix: A 2D slice of Vector (i.e., []Vector) used for batched operations and multi-dimensional data.
    • AttentionWeights: A type alias for Vector, used to provide semantic clarity in attention-related function signatures.
    type Vector []float64
type Matrix []Vector
type AttentionWeights = Vector

  3. Best practices for using go-attention

    main

    To ensure performance and reliability when building with go-attention, follow these guidelines:

    1. Use canonical functions: Always prefer DotProduct and DotProductAttention over deprecated variants.
    2. Check errors: Always handle the error return value from mathematical operations.
    3. Reuse memory pools: For high-frequency operations, use GetVectorFromPool and PutVectorToPool to reduce allocation overhead.
    4. Validate dimensions: Use validateMatrixDimensions when performing operations involving multiple matrices.
    5. Configure appropriately: Tune ParallelConfig and PerformanceConfig based on your specific hardware and workload requirements.

  5. Configure and monitor performance

    main

    The library includes built-in performance monitoring and tuning via PerformanceConfig.

    Configuration Options:

    • EnableMonitoring (bool): Enables tracking of operation counts and timing.
    • EnableAutoTuning (bool): Enables automatic algorithm/parallelism tuning.
    • MinVectorSize (int): Minimum size for parallel vector operations.
    • MinMatrixSize (int): Minimum size for parallel matrix operations.
    • MaxWorkers (int): Maximum number of workers for parallel tasks.

    Key Functions:

    • SetPerformanceConfig(config PerformanceConfig): Updates global settings.
    • DefaultPerformanceConfig(): Returns a recommended default config.
    • GetPerformanceStats(operation string) (*PerformanceStats, bool): Retrieves stats for a specific operation.
    • GetAllPerformanceStats() map[string]*PerformanceStats: Returns all collected stats.
    • ResetPerformanceStats(): Clears all stats.
    config := attention.PerformanceConfig{
    EnableMonitoring: true,
    MaxWorkers:       4,
}
attention.SetPerformanceConfig(config)

// Later, check stats
if stats, ok := attention.GetPerformanceStats("DotProduct"); ok {
    fmt.Printf("Avg time: %v\n", stats.AverageTime)
}

  6. Use Basic Dot-Product Attention

    main

    The DotProductAttention function implements the simplest form of attention mechanism, useful for basic sequence processing. It takes a query vector, a keys matrix, and a values matrix, returning the weighted output and the attention weights used.

    import "github.com/takara-ai/go-attention/attention"

// Create query-key-value setup
query := attention.Vector{1.0, 0.0, 1.0, 0.0}  // Pattern to search for
keys := attention.Matrix{
    {1.0, 0.0, 1.0, 0.0},  // Similar to query
    {0.0, 1.0, 0.0, 1.0},  // Different from query
    {0.5, 0.5, 0.5, 0.5},  // Neutral pattern
}
values := attention.Matrix{
    {1.0, 2.0},  // Value for similar key
    {3.0, 4.0},  // Value for different key
    {5.0, 6.0},  // Value for neutral key
}

// Compute attention
output, weights, err := attention.DotProductAttention(query, keys, values)
if err != nil {
    log.Fatal(err)
}

  7. Use a Full Transformer Layer

    main

    A complete transformer layer includes self-attention, layer normalization, position-wise feed-forward networks, and residual connections. Use transformer.NewTransformerLayer with a TransformerConfig to create the layer, then call Forward on your input matrix.

    import (
    "github.com/takara-ai/go-attention/transformer"
    "github.com/takara-ai/go-attention/attention"
)

// Configure transformer layer
config := transformer.TransformerConfig{
    DModel:      64,       // Size of token embeddings
    NumHeads:    4,        // Number of attention heads
    DHidden:     256,      // Size of feed-forward hidden layer
    DropoutRate: 0.1,      // For regularization
}

// Create transformer layer
layer, err := transformer.NewTransformerLayer(config)
if err != nil {
    log.Fatal(err)
}

// Create input sequence [seq_len × d_model]
seqLen := 3
input := make(attention.Matrix, seqLen)
for i := range input {
    input[i] = make(attention.Vector, config.DModel)
    // Fill with your embedding data...
}

// Process through transformer
output, err := layer.Forward(input)
if err != nil {
    log.Fatal(err)
}

  8. Use Multi-Head Attention

    main

    Multi-Head Attention (MHA) captures different types of relationships in parallel. You must first configure it using MultiHeadConfig and initialize it with NewMultiHeadAttention. The Forward method processes batched input sequences.

    import "github.com/takara-ai/go-attention/attention"

// Configure multi-head attention
config := attention.MultiHeadConfig{
    NumHeads:    4,        // Number of parallel attention heads
    DModel:      64,       // Size of input/output embeddings
    DKey:        16,       // Size per head (DModel/NumHeads)
    DValue:      16,       // Size per head (DModel/NumHeads)
    DropoutRate: 0.1,      // For regularization
}

// Create the attention module
mha, err := attention.NewMultiHeadAttention(config)
if err != nil {
    log.Fatal(err)
}

// Process sequences (batched input)
batchSize, seqLen := 2, 3  // Process 2 sequences, each with 3 tokens

// Create input matrices [batchSize × seqLen × DModel]
queries := make(attention.Matrix, batchSize*seqLen)
keys := make(attention.Matrix, batchSize*seqLen)
values := make(attention.Matrix, batchSize*seqLen)

// Initialize your matrices with actual data...

// Process through multi-head attention
output, err := mha.Forward(queries, keys, values)
if err != nil {
    log.Fatal(err)
}

  9. Use Feed-Forward Networks in the transformer package

    main

    The FeedForward struct implements a standard feed-forward network.

    Initialization: Use NewFeedForward(dModel, dHidden int) where dModel is the input/output dimension and dHidden is the hidden layer dimension.

    Usage: Call ff.Forward(input attention.Matrix) where the input matrix has shape [seq_len, d_model]. It returns the processed attention.Matrix of the same shape.

    ff := transformer.NewFeedForward(512, 2048)
output, err := ff.Forward(inputMatrix)

  10. Initialize and run a complete Transformer Layer

    main

    A TransformerLayer combines self-attention and a feed-forward network. It is configured using TransformerConfig.

    TransformerConfig fields:

    • DModel (int): Size of token embeddings.
    • NumHeads (int): Number of attention heads.
    • DHidden (int): Size of feed-forward hidden layer.
    • DropoutRate (float64): Regularization factor.

    Workflow:

    1. Create with NewTransformerLayer(config).
    2. Run the forward pass with t.Forward(input attention.Matrix) where input is [seq_len, d_model].
    config := transformer.TransformerConfig{
    DModel:      512,
    NumHeads:    8,
    DHidden:     2048,
    DropoutRate: 0.1,
}

layer, err := transformer.NewTransformerLayer(config)
output, err := layer.Forward(inputMatrix)

  11. Compute Scaled Dot-Product Attention

    main

    Use DotProductAttention to compute the canonical scaled dot-product attention mechanism.

    Parameters:

    • query Vector: The query vector of dimension [d_k].
    • keys Matrix: The key matrix of dimension [n, d_k].
    • values Matrix: The value matrix of dimension [n, d_v].

    Returns:

    • Vector: The attended output of dimension [d_v].
    • AttentionWeights: The calculated weights of dimension [n].
    • error: Error if dimensions are incompatible.

    BestDotProductAttention is an available alias for this function.

    output, weights, err := attention.DotProductAttention(query, keys, values)

  12. Use optimized matrix multiplication

    main

    For high-performance requirements, use MatrixMultiplyOptimized(a, b Matrix) (Matrix, error). This function implements blocking and cache-friendly access patterns to maximize throughput compared to standard implementations.

    result, err := MatrixMultiplyOptimized(matrixA, matrixB)
if err != nil {
    // Handle error
}