Coverage for transformer_lens/utilities/activation_functions.py: 100%

1"""Activation Functions. 

2 

3Utilities for interacting with all supported activation functions. 

4""" 

5 

6from typing import Callable, Dict, Optional, cast 

7 

8import numpy as np 

9import torch 

10import torch.nn as nn 

11import torch.nn.functional as F 

12from jaxtyping import Float 

13 

14 

15def gelu_new( 

16 input: Float[torch.Tensor, "batch pos d_mlp"], 

17) -> Float[torch.Tensor, "batch pos d_mlp"]: 

18 # Implementation of GeLU used by GPT2 - subtly different from PyTorch's 

19 return ( 

20 0.5 

21 * input 

22 * (1.0 + torch.tanh(np.sqrt(2.0 / np.pi) * (input + 0.044715 * torch.pow(input, 3.0)))) 

23 ) 

24 

25 

26def gelu_fast( 

27 input: Float[torch.Tensor, "batch pos d_mlp"], 

28) -> Float[torch.Tensor, "batch pos d_mlp"]: 

29 return 0.5 * input * (1.0 + torch.tanh(input * 0.7978845608 * (1.0 + 0.044715 * input * input))) 

30 

31 

32def gelu_pytorch_tanh(input: torch.Tensor) -> torch.Tensor: 

33 """Approximation of the gelu activation function, used in some older models.""" 

34 return F.gelu(input, approximate="tanh") 

35 

36 

37def solu(input: Float[torch.Tensor, "batch pos d_mlp"]) -> Float[torch.Tensor, "batch pos d_mlp"]: 

38 """ 

39 SoLU activation function as described by 

40 https://transformer-circuits.pub/2022/solu/index.html. 

41 

42 LayerNorm implemented by the MLP class. 

43 """ 

44 return input * F.softmax(input, dim=-1) 

45 

46 

47class XIELU(nn.Module): 

48 """Trainable xIELU activation function. 

49 

50 See https://arxiv.org/abs/2411.13010 

51 

52 Matches HuggingFace's XIELUActivation parameterization: alpha_p and alpha_n 

53 are stored in softplus-inverse space, and beta is a non-trainable buffer. 

54 """ 

55 

56 def __init__( 

57 self, 

58 alpha_p_init: float = 0.8, 

59 alpha_n_init: float = 0.8, 

60 beta_init: float = 0.5, 

61 eps: float = -1e-6, 

62 ): 

63 super().__init__() 

64 # Store in softplus-inverse space to match HF's XIELUActivation 

65 self.alpha_p = nn.Parameter( 

66 torch.log(torch.expm1(torch.tensor(alpha_p_init, dtype=torch.float32))) 

67 ) 

68 self.alpha_n = nn.Parameter( 

69 torch.log(torch.expm1(torch.tensor(alpha_n_init - beta_init, dtype=torch.float32))) 

70 ) 

71 self.beta: torch.Tensor 

72 self.eps: torch.Tensor 

73 self.register_buffer("beta", torch.tensor(beta_init, dtype=torch.float32)) 

74 self.register_buffer("eps", torch.tensor(eps, dtype=torch.float32)) 

75 

76 def forward( 

77 self, input: Float[torch.Tensor, "batch pos d_mlp"] 

78 ) -> Float[torch.Tensor, "batch pos d_mlp"]: 

79 alpha_p = F.softplus(self.alpha_p) 

80 alpha_n = self.beta + F.softplus(self.alpha_n) 

81 return torch.where( 

82 input > 0, 

83 alpha_p * input * input + self.beta * input, 

84 (torch.expm1(torch.min(input, self.eps)) - input) * alpha_n + self.beta * input, 

85 ) 

86 

87 

88def xielu(input: Float[torch.Tensor, "batch pos d_mlp"]) -> Float[torch.Tensor, "batch pos d_mlp"]: 

89 """Fixed-parameter xIELU activation function as described by 

90 https://arxiv.org/abs/2411.13010 

91 

92 Original code: https://github.com/rubber-duck-debug/xielu 

93 

94 Uses default parameter values. For trainable parameters, use the XIELU class. 

95 """ 

96 alpha_p: float = 0.8 

97 alpha_n: float = 0.8 

98 beta: float = 0.5 

99 eps = torch.tensor(-1e-6) 

100 

101 return torch.where( 

102 input > 0, 

103 alpha_p * input * input + beta * input, 

104 (torch.expm1(torch.min(input, eps)) - input) * alpha_n + beta * input, 

105 ) 

106 

107 

108def relu2(input: Float[torch.Tensor, "batch pos d_mlp"]) -> Float[torch.Tensor, "batch pos d_mlp"]: 

109 """Squared-ReLU (ReLU^2) activation function: ``relu(x) ** 2``. 

110 

111 Used by Arcee (ArceeForCausalLM / AFM-4.5B) in an ungated MLP. Matches 

112 HuggingFace's ``ReLUSquaredActivation`` (registered as ``"relu2"`` in 

113 ``transformers`` ACT2FN), so the key here mirrors the HF ``hidden_act`` string. 

114 The squaring drives high activation sparsity, which is useful for neuron-level 

115 interpretability (https://arxiv.org/abs/2109.08668). 

116 """ 

117 relu_applied = F.relu(input) 

118 return relu_applied * relu_applied 

119 

120 

121SOFTCAP_DISABLED: float = -1.0 

122 

123 

124def softcap_enabled(cap: Optional[float]) -> bool: 

125 """A soft-cap is active only for a strictly positive value; -1.0 / 0 / None mean disabled.""" 

126 return cap is not None and cap > 0 

127 

128 

129def apply_softcap(x: torch.Tensor, cap: Optional[float]) -> torch.Tensor: 

130 """Gemma-style ``cap * tanh(x / cap)`` when enabled, identity otherwise.""" 

131 if not softcap_enabled(cap): 

132 return x 

133 cap = cast(float, cap) 

134 return cap * torch.tanh(x / cap) 

135 

136 

137# Convenient type for the format of each activation function 

138ActivationFunction = Callable[..., torch.Tensor] 

139 

140# All currently supported activation functions. To add a new function, simply 

141# put the name of the function as the key, and the value as the actual callable. 

142SUPPORTED_ACTIVATIONS: Dict[str, ActivationFunction] = { 

143 "solu": solu, 

144 "solu_ln": solu, 

145 "gelu_new": gelu_new, 

146 "gelu_fast": gelu_fast, 

147 "silu": F.silu, 

148 "relu": F.relu, 

149 "gelu": F.gelu, 

150 "gelu_pytorch_tanh": gelu_pytorch_tanh, 

151 "xielu": xielu, 

152 "relu2": relu2, 

153}